executor.py

from codesafety import CodeSafety

import os
import json
import subprocess
import datetime
import logging

from uuid import uuid4

logger = logging.getLogger('executor')

def gen_filename():
    return f"{uuid4()}-{datetime.datetime.now().strftime('%Y%m%d')}.py".replace('-', "_")

class Executor:
    """
    Executes arbitrary user-defined python functions by creating
    temporary script files within an `executions` directory
    and calling the generated file with `subprocess.run`.

    Class Attributes:
        EXEC_PATH : str
            directory that stores temporary user scripts during execution
            will be auto-created if it does not already exist
        TIMEOUT : int
            max time in wall-clock seconds a user's function will run
            within a subprocess
        STDOUT_SPLIT_MAGIC : str
            a magic string used to identify the user's function
            result and separate it within stdout
        DEFAULT_IMPORTS : list[str]
            a list of imports to add into the user script
    """

    EXEC_PATH = "executions"
    TIMEOUT   = 5 # seconds
    STDOUT_SPLIT_MAGIC = "||------*|*+=|=+*|*------||"

    DEFAULT_IMPORTS = [
        "numpy"
        , "pandas"
        , "os"
        , "json"
        , "re"
        , "datetime"
    ]

    @classmethod
    def setup(cls) -> None:
        """ Creates the execution directory """
        os.makedirs(cls.EXEC_PATH, exist_ok=True)

    @classmethod
    def import_headers(cls) -> str:
        """
        Reads DEFAULT_IMPORTS to create valid import statements
        for a downstream user script
        """
        return '\n'.join(
            f"import {library}"
            for library in cls.DEFAULT_IMPORTS
        ) + '\n'

    @classmethod
    def gen_script(cls, script : str) -> str:
        """
        Generates a temporary python script loaded with
        DEFAULT_IMPORTS, a helper function `handle_output`,
        and the user-defined `main` function

        CodeSafety will verify that the function isn't malicious
        (to a degree) and will raise an exception if the internal
        heuristics identify malicious behavior

        Creates a temporary file within `EXEC_PATH` and returns
        the filename
        """
        fname = gen_filename()

        CodeSafety.check_user_main_function(
            fname, script
        )

        imports = cls.import_headers()

        script_text = f"""
{imports}

{script}

def handle_output(res):
    match res:
        case str():
            print(json.dumps(json.loads(res)))
        case dict():
            print(json.dumps(res))
        case _:
            raise Exception(str(type(res)) + " isn't handled currently!")

if __name__ == "__main__":
    result = main()
    print("{cls.STDOUT_SPLIT_MAGIC}")
    print(handle_output(result))

"""

        with open(os.path.join(cls.EXEC_PATH, fname), 'w') as f:
            f.write(script_text)

        return fname

    @classmethod
    def cleanup(cls, fname : str) -> None:
        """
        Removes the temporary user script at {EXEC_PATH}/{fname}
        """
        if fname is None:
            return

        path = os.path.join(cls.EXEC_PATH, fname)

        if os.path.exists(path):
            os.remove(path)

    @classmethod
    def handle_execution_result(
        cls
        , result : subprocess.CompletedProcess | None = None
        , error  : Exception | None = None
    ) -> dict:
        """
        Handles the success or failure of a specific user script's
        execution. If `error` exists, the function raises it immediately,
        otherwise result will be populated.

        If result exists, parses the stdout from the subprocess
        to return a dictionary containing the user-defined
        script result along with the rest of stdout
        """
        if error is not None:
            raise error

        status = result.returncode

        if status != 0:
            raise Exception(f"Process exited with non-zero status: {status}!")

        raw_stdout = result.stdout.decode()
        splits = raw_stdout.split('\n')

        stdout = []
        output = None

        for i, line in enumerate(splits):
            if line.strip() == cls.STDOUT_SPLIT_MAGIC:
                output = json.loads(splits[i + 1])
                break
            else:
                stdout.append(line)

        return {
            'result': output
            , 'stdout': ''.join(stdout)
        }

    @classmethod
    def execute(cls, fname : str, nsjail : bool = False) -> dict:
        """
        Executes a given filename through subprocess
        and hands the result/error to handle_execution_result

        nsjail flag can be set to force execution in a jailed environment
        """
        error = result = None

        logger.info(f"Executing {fname}...")

        try:
            if nsjail:
                cmd = [
                    'nsjail', '-C', '/etc/nsjail.cfg'
                    , '--chroot', '/', '--'
                    , '/usr/bin/python3'
                    , os.path.join('/app', cls.EXEC_PATH, fname)
                ]
            else:
                cmd = ['python3', os.path.join(cls.EXEC_PATH, fname)]

            result = subprocess.run(
                cmd
                , timeout=cls.TIMEOUT
                , stdout=subprocess.PIPE
            )
        except Exception as e:
            logger.warning(e)
            error = e

        return cls.handle_execution_result(
            result=result, error=error
        )

    @classmethod
    def run(cls, script : str, nsjail : bool = False) -> dict:
        """
        The main entry point of Executor, runs the entire
        setup -> verify -> generate -> execute -> cleanup
        loop for a given user-defined script
        """
        cls.setup()

        result = None
        fname = None

        try:
            fname  = cls.gen_script(script)
            result = cls.execute(fname, nsjail=nsjail)
        except Exception as e:
            # TODO: log errors to db or inspect results
            result = {'error': e}
        finally:
            cls.cleanup(fname)

        return result