Merge branch 'main' of github.com:facundobatista/pyempaq into main

pyempaq/unpacker.py

@@ -12,6 +12,8 @@
 import time
 import venv
 import zipfile
+from typing import List, Dict
+
 
 # this is the directory for the NEW virtualenv created for the project (not the packed
 # one to run unpacker itself)
@@ -52,62 +54,77 @@ def get_python_exec(project_dir):
     return executable
 
 
-log("Pyempaq start")
-
-# parse pyempaq metadata from the zip file
-pyempaq_filepath = pathlib.Path.cwd() / sys.argv[0]
-zf = zipfile.ZipFile(pyempaq_filepath)
-metadata = json.loads(zf.read("metadata.json").decode("utf8"))
-log("Loaded metadata: {}", metadata)
-
-# load appdirs from the builtin venv
-sys.path.insert(0, f"{pyempaq_filepath}/venv/")
-import appdirs  # NOQA: this is an import not at top of file because paths needed to be fixed
-
-pyempaq_dir = pathlib.Path(appdirs.user_data_dir()) / 'pyempaq'
-pyempaq_dir.mkdir(parents=True, exist_ok=True)
-log("Temp base dir: {!r}", str(pyempaq_dir))
-
-# create a temp dir and extract the project there
-timestamp = time.strftime("%Y%m%d%H%M%S", time.gmtime(pyempaq_filepath.stat().st_ctime))
-project_dir = pyempaq_dir / "{}-{}".format(metadata["project_name"], timestamp)
-original_project_dir = project_dir / "orig"
-if project_dir.exists():
-    log("Reusing project dir {!r}", str(project_dir))
-else:
-    log("Creating project dir {!r}", str(project_dir))
-    project_dir.mkdir()
-
-    log("Extracting pyempaq content")
-    zf.extractall(path=project_dir)
-
-    venv_requirements = metadata["requirement_files"]
-    if venv_requirements:
-        log("Creating payload virtualenv")
-        venv_dir = project_dir / PROJECT_VENV_DIR
-        venv.create(venv_dir, with_pip=True)
-        pip_exec = find_venv_bin(venv_dir, "pip3")
-        cmd = [str(pip_exec), "install"]
-        for req_file in venv_requirements:
-            cmd += ["-r", str(original_project_dir / req_file)]
-        log("Installing dependencies: {}", cmd)
-        subprocess.run(cmd, check=True)
-        log("Virtualenv setup finished")
+def build_command(python_exec: str, metadata: Dict[str, str], sys_args: List[str]) -> List[str]:
+    """Build the command to be executed."""
+    if metadata["exec_style"] == "script":
+        cmd = [python_exec, metadata["exec_value"]]
+    elif metadata["exec_style"] == "module":
+        cmd = [python_exec, "-m", metadata["exec_value"]]
+    elif metadata["exec_style"] == "entrypoint":
+        cmd = [python_exec] + metadata["exec_value"]
 
+    if sys_args:
+        cmd.extend(sys_args)
+    else:
+        cmd.extend(metadata["exec_default_args"])
+    return cmd
+
+
+def run():
+    """Run the unpacker."""
+    log("Pyempaq start")
+
+    # parse pyempaq metadata from the zip file
+    pyempaq_filepath = pathlib.Path.cwd() / sys.argv[0]
+    zf = zipfile.ZipFile(pyempaq_filepath)
+    metadata = json.loads(zf.read("metadata.json").decode("utf8"))
+    log("Loaded metadata: {}", metadata)
+
+    # load appdirs from the builtin venv
+    sys.path.insert(0, f"{pyempaq_filepath}/venv/")
+    import appdirs  # NOQA: this is an import not at top of file because paths needed to be fixed
+
+    pyempaq_dir = pathlib.Path(appdirs.user_data_dir()) / 'pyempaq'
+    pyempaq_dir.mkdir(parents=True, exist_ok=True)
+    log("Temp base dir: {!r}", str(pyempaq_dir))
+
+    # create a temp dir and extract the project there
+    timestamp = time.strftime("%Y%m%d%H%M%S", time.gmtime(pyempaq_filepath.stat().st_ctime))
+    project_dir = pyempaq_dir / "{}-{}".format(metadata["project_name"], timestamp)
+    original_project_dir = project_dir / "orig"
+    if project_dir.exists():
+        log("Reusing project dir {!r}", str(project_dir))
     else:
-        log("Skipping virtualenv (no requirements)")
-
-python_exec = str(get_python_exec(project_dir))
-os.chdir(original_project_dir)
-
-if metadata["exec_style"] == "script":
-    cmd = [python_exec, metadata["exec_value"]]
-elif metadata["exec_style"] == "module":
-    cmd = [python_exec, "-m", metadata["exec_value"]]
-elif metadata["exec_style"] == "entrypoint":
-    cmd = [python_exec] + metadata["exec_value"]
-cmd.extend(metadata["exec_default_args"])
-
-log("Running payload: {}", cmd)
-subprocess.run(cmd)
-log("Pyempaq done")
+        log("Creating project dir {!r}", str(project_dir))
+        project_dir.mkdir()
+
+        log("Extracting pyempaq content")
+        zf.extractall(path=project_dir)
+
+        venv_requirements = metadata["requirement_files"]
+        if venv_requirements:
+            log("Creating payload virtualenv")
+            venv_dir = project_dir / PROJECT_VENV_DIR
+            venv.create(venv_dir, with_pip=True)
+            pip_exec = find_venv_bin(venv_dir, "pip3")
+            cmd = [str(pip_exec), "install"]
+            for req_file in venv_requirements:
+                cmd += ["-r", str(original_project_dir / req_file)]
+            log("Installing dependencies: {}", cmd)
+            subprocess.run(cmd, check=True)
+            log("Virtualenv setup finished")
+
+        else:
+            log("Skipping virtualenv (no requirements)")
+
+    python_exec = str(get_python_exec(project_dir))
+    os.chdir(original_project_dir)
+
+    cmd = build_command(python_exec, metadata, sys.argv[1:])
+    log("Running payload: {}", cmd)
+    subprocess.run(cmd)
+    log("Pyempaq done")
+
+
+if __name__ == "__main__":
+    run()

tests/test_unpacker.py

@@ -0,0 +1,62 @@
+# Copyright 2021 Facundo Batista
+# Licensed under the GPL v3 License
+# For further info, check https://github.com/facundobatista/pyempaq
+
+"""Unpacker tests."""
+
+from pyempaq.unpacker import build_command
+
+
+def test_buildcommand_script_default_empty():
+    """Build a command in "script" mode."""
+    metadata = {
+        "exec_style": "script",
+        "exec_value": "mystuff.py",
+        "exec_default_args": [],
+    }
+    cmd = build_command("python.exe", metadata, [])
+    assert cmd == ["python.exe", "mystuff.py"]
+
+
+def test_buildcommand_module_default_empty():
+    """Build a command in "module" mode."""
+    metadata = {
+        "exec_style": "module",
+        "exec_value": "mymodule",
+        "exec_default_args": [],
+    }
+    cmd = build_command("python.exe", metadata, [])
+    assert cmd == ["python.exe", "-m", "mymodule"]
+
+
+def test_buildcommand_entrypoint_default_empty():
+    """Build a command in "entrypoint" mode."""
+    metadata = {
+        "exec_style": "entrypoint",
+        "exec_value": ["whatever", "you", "want"],
+        "exec_default_args": [],
+    }
+    cmd = build_command("python.exe", metadata, [])
+    assert cmd == ["python.exe", "whatever", "you", "want"]
+
+
+def test_buildcommand_script_default_nonempty():
+    """Build a command with default args."""
+    metadata = {
+        "exec_style": "script",
+        "exec_value": "mystuff.py",
+        "exec_default_args": ["--foo", "3"],
+    }
+    cmd = build_command("python.exe", metadata, [])
+    assert cmd == ["python.exe", "mystuff.py", "--foo", "3"]
+
+
+def test_buildcommand_script_sysargs():
+    """Build a command with user passed args."""
+    metadata = {
+        "exec_style": "script",
+        "exec_value": "mystuff.py",
+        "exec_default_args": ["--foo", "3"],
+    }
+    cmd = build_command("python.exe", metadata, ["--bar"])
+    assert cmd == ["python.exe", "mystuff.py", "--bar"]