#include <Python.h>

#include <marshal.h>

#include "pycore_fileutils.h" // _Py_stat_struct

#include <pycore_import.h>

#include <stdio.h>

#include <stdlib.h> // malloc()

#include <sys/types.h>

#include <sys/stat.h>

#ifndef MS_WINDOWS

# include <unistd.h>

static const struct _frozen no_modules[] = {

{0, 0, 0} /* sentinel */

};

static const struct _module_alias aliases[] = {

{0, 0} /* sentinel */

};

const struct _frozen *_PyImport_FrozenBootstrap;

const struct _frozen *_PyImport_FrozenStdlib;

const struct _frozen *_PyImport_FrozenTest;

const struct _frozen *PyImport_FrozenModules;

const struct _module_alias *_PyImport_FrozenAliases;

static const char header[] =

"/* Auto-generated by Programs/_freeze_module.c */";

static void

runtime_init(void)

{

PyInitConfig *config = PyInitConfig_Create();

if (config == NULL) {

printf("memory allocation failed\n");

exit(1);

}

if (PyInitConfig_SetInt(config, "site_import", 0) < 0) {

goto error;

}

if (PyInitConfig_SetStr(config, "program_name", "./_freeze_module") < 0) {

goto error;

}

/* Don't install importlib, since it could execute outdated bytecode. */

if (PyInitConfig_SetInt(config, "_install_importlib", 0) < 0) {

goto error;

}

if (PyInitConfig_SetInt(config, "_init_main", 0) < 0) {

goto error;

}

if (Py_InitializeFromInitConfig(config) < 0) {

goto error;

}

PyInitConfig_Free(config);

return;

error:

{

const char *err_msg;

(void)PyInitConfig_GetError(config, &err_msg);

printf("Python init error: %s\n", err_msg);

PyInitConfig_Free(config);

exit(1);

}

}

static const char *

read_text(const char *inpath)

{

FILE *infile = fopen(inpath, "rb");

if (infile == NULL) {

fprintf(stderr, "cannot open '%s' for reading\n", inpath);

return NULL;

}

struct _Py_stat_struct stat;

if (_Py_fstat_noraise(fileno(infile), &stat)) {

fprintf(stderr, "cannot fstat '%s'\n", inpath);

fclose(infile);

return NULL;

}

size_t text_size = (size_t)stat.st_size;

char *text = (char *) malloc(text_size + 1);

if (text == NULL) {

fprintf(stderr, "could not allocate %ld bytes\n", (long) text_size);

fclose(infile);

return NULL;

}

size_t n = fread(text, 1, text_size, infile);

fclose(infile);

if (n < text_size) {

fprintf(stderr, "read too short: got %ld instead of %ld bytes\n",

(long) n, (long) text_size);

free(text);

return NULL;

}

text[text_size] = '\0';

return (const char *)text;

}

static PyObject *

compile_and_marshal(const char *name, const char *text)

{

char *filename = (char *) malloc(strlen(name) + 10);

if (filename == NULL) {

return PyErr_NoMemory();

}

sprintf(filename, "<frozen %s>", name);

PyObject *code = Py_CompileStringExFlags(text, filename,

Py_file_input, NULL, 0);

free(filename);

if (code == NULL) {

return NULL;

}

assert(Py_MARSHAL_VERSION >= 6);

PyObject *marshalled = PyMarshal_WriteObjectToString(code, Py_MARSHAL_VERSION);

Py_CLEAR(code);

if (marshalled == NULL) {

return NULL;

}

assert(PyBytes_CheckExact(marshalled));

return marshalled;

}

static char *

get_varname(const char *name, const char *prefix)

{

size_t n = strlen(prefix);

char *varname = (char *) malloc(strlen(name) + n + 1);

if (varname == NULL) {

return NULL;

}

(void)strcpy(varname, prefix);

for (size_t i = 0; name[i] != '\0'; i++) {

if (name[i] == '.') {

varname[n++] = '_';

}

else {

varname[n++] = name[i];

}

}

varname[n] = '\0';

return varname;

}

static void

write_code(FILE *outfile, PyObject *marshalled, const char *varname)

{

unsigned char *data = (unsigned char *) PyBytes_AS_STRING(marshalled);

size_t data_size = PyBytes_GET_SIZE(marshalled);

fprintf(outfile, "const unsigned char %s[] = {\n", varname);

for (size_t n = 0; n < data_size; n += 16) {

size_t i, end = Py_MIN(n + 16, data_size);

fprintf(outfile, " ");

for (i = n; i < end; i++) {

fprintf(outfile, "%u,", (unsigned int) data[i]);

}

fprintf(outfile, "\n");

}

fprintf(outfile, "};\n");

}

static int

write_frozen(const char *outpath, const char *inpath, const char *name,

PyObject *marshalled)

{

/* Open the file in text mode. The hg checkout should be using the eol extension,

FILE *outfile = fopen(outpath, "w");

if (outfile == NULL) {

fprintf(stderr, "cannot open '%s' for writing\n", outpath);

return -1;

}

fprintf(outfile, "%s\n", header);

char *arrayname = get_varname(name, "_Py_M__");

if (arrayname == NULL) {

fprintf(stderr, "memory error: could not allocate varname\n");

fclose(outfile);

return -1;

}

write_code(outfile, marshalled, arrayname);

free(arrayname);

if (ferror(outfile)) {

fprintf(stderr, "error when writing to '%s'\n", outpath);

fclose(outfile);

return -1;

}

fclose(outfile);

return 0;

}

main(int argc, char *argv[])

{

const char *name, *inpath, *outpath;

_PyImport_FrozenBootstrap = no_modules;

_PyImport_FrozenStdlib = no_modules;

_PyImport_FrozenTest = no_modules;

PyImport_FrozenModules = NULL;

_PyImport_FrozenAliases = aliases;

if (argc != 4) {

fprintf(stderr, "need to specify the name, input and output paths\n");

return 2;

}

name = argv[1];

inpath = argv[2];

outpath = argv[3];

runtime_init();

const char *text = read_text(inpath);

if (text == NULL) {

goto error;

}

PyObject *marshalled = compile_and_marshal(name, text);

free((char *)text);

if (marshalled == NULL) {

goto error;

}

int res = write_frozen(outpath, inpath, name, marshalled);

Py_DECREF(marshalled);

if (res != 0) {

goto error;

}

Py_Finalize();

return 0;

error:

PyErr_Print();

Py_Finalize();

return 1;

}