#include "Python.h"

#include "structmember.h"

#include "osdefs.h"

#include "marshal.h"

#include <time.h>

#define IS_SOURCE   0x0

#define IS_BYTECODE 0x1

#define IS_PACKAGE  0x2

struct st_zip_searchorder {

	char suffix[14];

	int type;

};

/* zip_searchorder defines how we search for a module in the Zip

   archive: we first search for a package __init__, then for

   non-package .pyc, .pyo and .py entries. The .pyc and .pyo entries

   are swapped by initzipimport() if we run in optimized mode. Also,

   '/' is replaced by SEP there. */

static struct st_zip_searchorder zip_searchorder[] = {

	{"/__init__.pyc", IS_PACKAGE | IS_BYTECODE},

	{"/__init__.pyo", IS_PACKAGE | IS_BYTECODE},

	{"/__init__.py", IS_PACKAGE | IS_SOURCE},

	{".pyc", IS_BYTECODE},

	{".pyo", IS_BYTECODE},

	{".py", IS_SOURCE},

	{"", 0}

};

/* zipimporter object definition and support */

typedef struct _zipimporter ZipImporter;

struct _zipimporter {

	PyObject_HEAD

	PyObject *archive;  /* pathname of the Zip archive */

	PyObject *prefix;   /* file prefix: "a/sub/directory/" */

	PyObject *files;    /* dict with file info {path: toc_entry} */

};

static PyObject *ZipImportError;

static PyObject *zip_directory_cache = NULL;

/* forward decls */

static PyObject *read_directory(char *archive);

static PyObject *get_data(char *archive, PyObject *toc_entry);

static PyObject *get_module_code(ZipImporter *self, char *fullname,

				 int *p_ispackage, char **p_modpath);

#define ZipImporter_Check(op) PyObject_TypeCheck(op, &ZipImporter_Type)

/* zipimporter.__init__

   Split the "subdirectory" from the Zip archive path, lookup a matching

   entry in sys.path_importer_cache, fetch the file directory from there

   if found, or else read it from the archive. */

static int

zipimporter_init(ZipImporter *self, PyObject *args, PyObject *kwds)

{

	char *path, *p, *prefix, buf[MAXPATHLEN+2];

	size_t len;

	if (!_PyArg_NoKeywords("zipimporter()", kwds))

		return -1;

	if (!PyArg_ParseTuple(args, "s:zipimporter",

			      &path))

		return -1;

	len = strlen(path);

	if (len == 0) {

		PyErr_SetString(ZipImportError, "archive path is empty");

		return -1;

	}

	if (len >= MAXPATHLEN) {

		PyErr_SetString(ZipImportError,

				"archive path too long");

		return -1;

	}

	strcpy(buf, path);

#ifdef ALTSEP

	for (p = buf; *p; p++) {

		if (*p == ALTSEP)

			*p = SEP;

	}

#endif

	path = NULL;

	prefix = NULL;

	for (;;) {

#ifndef RISCOS

		struct stat statbuf;

		int rv;

		rv = stat(buf, &statbuf);

		if (rv == 0) {

			/* it exists */

			if (S_ISREG(statbuf.st_mode))

				/* it's a file */

				path = buf;

			break;

		}

#else

		if (object_exists(buf)) {

			/* it exists */

			if (isfile(buf))

				/* it's a file */

				path = buf;

			break;

		}

#endif

		/* back up one path element */

		p = strrchr(buf, SEP);

		if (prefix != NULL)

			*prefix = SEP;

		if (p == NULL)

			break;

		*p = '\0';

		prefix = p;

	}

	if (path != NULL) {

		PyObject *files;

		files = PyDict_GetItemString(zip_directory_cache, path);

		if (files == NULL) {

			files = read_directory(buf);

			if (files == NULL)

				return -1;

			if (PyDict_SetItemString(zip_directory_cache, path,

						 files) != 0)

				return -1;

		}

		else

			Py_INCREF(files);

		self->files = files;

	}

	else {

		PyErr_SetString(ZipImportError, "not a Zip file");

		return -1;

	}

	if (prefix == NULL)

		prefix = "";

	else {

		prefix++;

		len = strlen(prefix);

		if (prefix[len-1] != SEP) {

			/* add trailing SEP */

			prefix[len] = SEP;

			prefix[len + 1] = '\0';

		}

	}

	self->archive = PyString_FromString(buf);

	if (self->archive == NULL)

		return -1;

	self->prefix = PyString_FromString(prefix);

	if (self->prefix == NULL)

		return -1;

	return 0;

}

/* GC support. */

static int

zipimporter_traverse(PyObject *obj, visitproc visit, void *arg)

{

	ZipImporter *self = (ZipImporter *)obj;

	Py_VISIT(self->files);

	return 0;

}

static void

zipimporter_dealloc(ZipImporter *self)

{

	PyObject_GC_UnTrack(self);

	Py_XDECREF(self->archive);

	Py_XDECREF(self->prefix);

	Py_XDECREF(self->files);

	Py_TYPE(self)->tp_free((PyObject *)self);

}

static PyObject *

zipimporter_repr(ZipImporter *self)

{

	char buf[500];

	char *archive = "???";

	char *prefix = "";

	if (self->archive != NULL && PyString_Check(self->archive))

		archive = PyString_AsString(self->archive);

	if (self->prefix != NULL && PyString_Check(self->prefix))

		prefix = PyString_AsString(self->prefix);

	if (prefix != NULL && *prefix)

		PyOS_snprintf(buf, sizeof(buf),

			      "<zipimporter object \"%.300s%c%.150s\">",

			      archive, SEP, prefix);

	else

		PyOS_snprintf(buf, sizeof(buf),

			      "<zipimporter object \"%.300s\">",

			      archive);

	return PyString_FromString(buf);

}

/* return fullname.split(".")[-1] */

static char *

get_subname(char *fullname)

{

	char *subname = strrchr(fullname, '.');

	if (subname == NULL)

		subname = fullname;

	else

		subname++;

	return subname;

}

/* Given a (sub)modulename, write the potential file path in the

   archive (without extension) to the path buffer. Return the

   length of the resulting string. */

static int

make_filename(char *prefix, char *name, char *path)

{

	size_t len;

	char *p;

	len = strlen(prefix);

	/* self.prefix + name [+ SEP + "__init__"] + ".py[co]" */

	if (len + strlen(name) + 13 >= MAXPATHLEN) {

		PyErr_SetString(ZipImportError, "path too long");

		return -1;

	}

	strcpy(path, prefix);

	strcpy(path + len, name);

	for (p = path + len; *p; p++) {

		if (*p == '.')

			*p = SEP;

	}

	len += strlen(name);

	assert(len < INT_MAX);

	return (int)len;

}

enum zi_module_info {

	MI_ERROR,

	MI_NOT_FOUND,

	MI_MODULE,

	MI_PACKAGE

};

/* Return some information about a module. */

static enum zi_module_info

get_module_info(ZipImporter *self, char *fullname)

{

	char *subname, path[MAXPATHLEN + 1];

	int len;

	struct st_zip_searchorder *zso;

	subname = get_subname(fullname);

	len = make_filename(PyString_AsString(self->prefix), subname, path);

	if (len < 0)

		return MI_ERROR;

	for (zso = zip_searchorder; *zso->suffix; zso++) {

		strcpy(path + len, zso->suffix);

		if (PyDict_GetItemString(self->files, path) != NULL) {

			if (zso->type & IS_PACKAGE)

				return MI_PACKAGE;

			else

				return MI_MODULE;

		}

	}

	return MI_NOT_FOUND;

}

/* Check whether we can satisfy the import of the module named by

   'fullname'. Return self if we can, None if we can't. */

static PyObject *

zipimporter_find_module(PyObject *obj, PyObject *args)

{

	ZipImporter *self = (ZipImporter *)obj;

	PyObject *path = NULL;

	char *fullname;

	enum zi_module_info mi;

	if (!PyArg_ParseTuple(args, "s|O:zipimporter.find_module",

			      &fullname, &path))

		return NULL;

	mi = get_module_info(self, fullname);

	if (mi == MI_ERROR)

		return NULL;

	if (mi == MI_NOT_FOUND) {

		Py_INCREF(Py_None);

		return Py_None;

	}

	Py_INCREF(self);

	return (PyObject *)self;

}

/* Load and return the module named by 'fullname'. */

static PyObject *

zipimporter_load_module(PyObject *obj, PyObject *args)

{

	ZipImporter *self = (ZipImporter *)obj;

	PyObject *code, *mod, *dict;

	char *fullname, *modpath;

	int ispackage;

	if (!PyArg_ParseTuple(args, "s:zipimporter.load_module",

			      &fullname))

		return NULL;

	code = get_module_code(self, fullname, &ispackage, &modpath);

	if (code == NULL)

		return NULL;

	mod = PyImport_AddModule(fullname);

	if (mod == NULL) {

		Py_DECREF(code);

		return NULL;

	}

	dict = PyModule_GetDict(mod);

	/* mod.__loader__ = self */

	if (PyDict_SetItemString(dict, "__loader__", (PyObject *)self) != 0)

		goto error;

	if (ispackage) {

		/* add __path__ to the module *before* the code gets

		   executed */

		PyObject *pkgpath, *fullpath;

		char *prefix = PyString_AsString(self->prefix);

		char *subname = get_subname(fullname);

		int err;

		fullpath = PyString_FromFormat("%s%c%s%s",

					PyString_AsString(self->archive),

					SEP,

					*prefix ? prefix : "",

					subname);

		if (fullpath == NULL)

			goto error;

		pkgpath = Py_BuildValue("[O]", fullpath);

		Py_DECREF(fullpath);

		if (pkgpath == NULL)

			goto error;

		err = PyDict_SetItemString(dict, "__path__", pkgpath);

		Py_DECREF(pkgpath);

		if (err != 0)

			goto error;

	}

	mod = PyImport_ExecCodeModuleEx(fullname, code, modpath);

	Py_DECREF(code);

	if (Py_VerboseFlag)

		PySys_WriteStderr("import %s # loaded from Zip %s\n",

				  fullname, modpath);

	return mod;

error:

	Py_DECREF(code);

	Py_DECREF(mod);

	return NULL;

}

/* Return a bool signifying whether the module is a package or not. */

static PyObject *

zipimporter_is_package(PyObject *obj, PyObject *args)

{

	ZipImporter *self = (ZipImporter *)obj;

	char *fullname;

	enum zi_module_info mi;

	if (!PyArg_ParseTuple(args, "s:zipimporter.is_package",

			      &fullname))

		return NULL;

	mi = get_module_info(self, fullname);

	if (mi == MI_ERROR)

		return NULL;

	if (mi == MI_NOT_FOUND) {

		PyErr_Format(ZipImportError, "can't find module '%.200s'",

			     fullname);

		return NULL;

	}

	return PyBool_FromLong(mi == MI_PACKAGE);

}

static PyObject *

zipimporter_get_data(PyObject *obj, PyObject *args)

{

	ZipImporter *self = (ZipImporter *)obj;

	char *path;

#ifdef ALTSEP

	char *p, buf[MAXPATHLEN + 1];

#endif

	PyObject *toc_entry;

	Py_ssize_t len;

	if (!PyArg_ParseTuple(args, "s:zipimporter.get_data", &path))

		return NULL;

#ifdef ALTSEP

	if (strlen(path) >= MAXPATHLEN) {

		PyErr_SetString(ZipImportError, "path too long");

		return NULL;

	}

	strcpy(buf, path);

	for (p = buf; *p; p++) {

		if (*p == ALTSEP)

			*p = SEP;

	}

	path = buf;

#endif

	len = PyString_Size(self->archive);

	if ((size_t)len < strlen(path) &&

	    strncmp(path, PyString_AsString(self->archive), len) == 0 &&

	    path[len] == SEP) {

		path = path + len + 1;

	}

	toc_entry = PyDict_GetItemString(self->files, path);

	if (toc_entry == NULL) {

		PyErr_SetFromErrnoWithFilename(PyExc_IOError, path);

		return NULL;

	}

	return get_data(PyString_AsString(self->archive), toc_entry);

}

static PyObject *

zipimporter_get_code(PyObject *obj, PyObject *args)

{

	ZipImporter *self = (ZipImporter *)obj;

	char *fullname;

	if (!PyArg_ParseTuple(args, "s:zipimporter.get_code", &fullname))

		return NULL;

	return get_module_code(self, fullname, NULL, NULL);

}

static PyObject *

zipimporter_get_source(PyObject *obj, PyObject *args)

{

	ZipImporter *self = (ZipImporter *)obj;

	PyObject *toc_entry;

	char *fullname, *subname, path[MAXPATHLEN+1];

	int len;

	enum zi_module_info mi;

	if (!PyArg_ParseTuple(args, "s:zipimporter.get_source", &fullname))

		return NULL;

	mi = get_module_info(self, fullname);

	if (mi == MI_ERROR)

		return NULL;

	if (mi == MI_NOT_FOUND) {

		PyErr_Format(ZipImportError, "can't find module '%.200s'",

			     fullname);

		return NULL;

	}

	subname = get_subname(fullname);

	len = make_filename(PyString_AsString(self->prefix), subname, path);

	if (len < 0)

		return NULL;

	if (mi == MI_PACKAGE) {

		path[len] = SEP;

		strcpy(path + len + 1, "__init__.py");

	}

	else

		strcpy(path + len, ".py");

	toc_entry = PyDict_GetItemString(self->files, path);

	if (toc_entry != NULL)

		return get_data(PyString_AsString(self->archive), toc_entry);

	/* we have the module, but no source */

	Py_INCREF(Py_None);

	return Py_None;

}

PyDoc_STRVAR(doc_find_module,

"find_module(fullname, path=None) -> self or None.\n\

\n\

Search for a module specified by 'fullname'. 'fullname' must be the\n\

fully qualified (dotted) module name. It returns the zipimporter\n\

instance itself if the module was found, or None if it wasn't.\n\

The optional 'path' argument is ignored -- it's there for compatibility\n\

with the importer protocol.");

PyDoc_STRVAR(doc_load_module,

"load_module(fullname) -> module.\n\

\n\

Load the module specified by 'fullname'. 'fullname' must be the\n\

fully qualified (dotted) module name. It returns the imported\n\

module, or raises ZipImportError if it wasn't found.");

PyDoc_STRVAR(doc_get_data,

"get_data(pathname) -> string with file data.\n\

\n\

Return the data associated with 'pathname'. Raise IOError if\n\

the file wasn't found.");

PyDoc_STRVAR(doc_is_package,

"is_package(fullname) -> bool.\n\

\n\

Return True if the module specified by fullname is a package.\n\

Raise ZipImportError is the module couldn't be found.");

PyDoc_STRVAR(doc_get_code,

"get_code(fullname) -> code object.\n\

\n\

Return the code object for the specified module. Raise ZipImportError\n\

is the module couldn't be found.");

PyDoc_STRVAR(doc_get_source,

"get_source(fullname) -> source string.\n\

\n\

Return the source code for the specified module. Raise ZipImportError\n\

is the module couldn't be found, return None if the archive does\n\

contain the module, but has no source for it.");