run.py

import argparse
import os.path
import subprocess

# Creates and runs a Docker image to analyze malware
MODES = ["build", "run"]
IMAGE_NAME = "malwarescanner"


def is_url(path: str) -> bool:
    """Checks if a given path is a URL"""
    return "http://" in path or "https://" in path


def main():
    parser = argparse.ArgumentParser(prog="Malware Scanner Runner",
                                     description="Creates and runs a Docker image to analyze malware")

    subparsers = parser.add_subparsers(dest="subcommand")
    subparsers.required = True

    # subparser for args to build the Docker image
    parser_build = subparsers.add_parser("build")

    # subparser for args to run the scanner
    parser_run = subparsers.add_parser("run")
    parser_run.add_argument('filename', type=str, help="The file path or URL for the malware sample")
    parser_run.add_argument('-p', '--password', dest='password', default=" ", required=False,
                            help="Password for password-protected archives")
    parser_run.add_argument('-l', "--level", choices=["debug", "info", "quiet"], default="info", required=False,
                            help="Choose the level of information the container will print")

    args = parser.parse_args()

    if args.subcommand == "build":
        result = subprocess.run(["docker", "build", "-t", IMAGE_NAME, "."])

    elif args.subcommand == "run":
        run_args = ["docker", "run", "-it", "--rm", "--name=scanner"]
        fp = args.filename
        pw = args.password
        level = args.level

        if is_url(fp):               # url
            run_args += [IMAGE_NAME, fp, pw, level]
            result = subprocess.run(run_args)
        elif os.path.exists(fp):
            # Checks to see if it should use a relative file path or absolute one
            # no promises this works on Windows
            if fp[0] == '~' or fp[0] == '/':
                real_fp = fp
            else:
                real_fp = os.getcwd() + "/" + fp
            # print(real_fp)
            if os.path.isfile(fp):     # file
                # TODO: relative file path entry instead of writing the whole path
                # can probably do some Python magic to get the pwd and attach fp
                file_name = fp.split('/')[-1]
                run_args += ["-v", real_fp[:-len(file_name)] + ":/app/mount", IMAGE_NAME, "/app/mount/" + file_name,
                             pw, level]
            elif os.path.isdir(fp):      # directory
                run_args += ["-v", real_fp + ":/app/mount", IMAGE_NAME, "/app/mount", pw, level]
            else:
                print("This should never happen")
            # print(run_args)
            result = subprocess.run(run_args)
        else:
            print("ERROR: " + str(args.filename) + " not found!")


if __name__ == "__main__":
    main()