import os
import re
import subprocess

config_path = os.getenv('ANCHORE_FRESHCLAM_CONFIG_FILE', '/home/anchore/clamav/freshclam.conf')
dbdir = os.getenv('ANCHORE_CLAMAV_DB_DIR', '/home/anchore/clamav/db')

CLAMAV_CMD = "clamscan --suppress-ok-results --infected --recursive --allmatch --archive-verbose --tempdir={tempdir} --database={database}"
REFRESH_CMD = "freshclam --stdout --datadir={database} --config-file={configfile}"
CLAMAV_FINDING_EXITCODE = 1
CLAMAV_NO_FINDING_EXITCODE = 0


class ScanFailedException(Exception):
    pass


class RefreshFailedException(Exception):
    pass


def run_clamav(path_to_scan: str, tempdir: str, dbdir: str) -> str:
    """
    Run the scan

    Example expected output for a finding from a tar:

    "t1.tar.gz: Unix.Trojan.MSShellcode-40 FOUND
    t1.tar.gz!POSIX_TAR:8615aed40f9faebd56f4dfd6af10e576e4fdc49af4f4507dd738c51ae18a7a55/layer.tar!POSIX_TAR:bin/busybox!...!(8)POSIX_TAR:elf_payload1: Unix.Trojan.MSShellcode-40 FOUND"

    :param path_to_scan: str path of file to scan
    :param tempdir: str dir path to use as tmp during scan
    :param dbdir: str path to clamav virus db
    :return:
    """

    cmd = CLAMAV_CMD.format(tempdir=tempdir, database=dbdir) + " " + path_to_scan
    print('Running command {}'.format(cmd))
    status, output = subprocess.getstatusoutput(cmd)
    if status in [CLAMAV_FINDING_EXITCODE, CLAMAV_NO_FINDING_EXITCODE]:
        return output
    else:
        raise ScanFailedException('clamav execution failed. returned exit code {}. Output = {}'.format(status, output))


def refresh_clamavdb(databasedir: str, configfile: str) -> str:
    """
    Run the refresh and return the output str

    :param databasedir: str path to where the clamav db resides
    :param configfile: str path to the config file
    :return: output of the command execution as a str
    """

    print('Refreshing ClamAV DB')
    cmd = REFRESH_CMD.format(database=databasedir, configfile=configfile)
    print('Running command {}'.format(cmd))
    status, output = subprocess.getstatusoutput(cmd)
    if status == 0:
        return output
    else:
        # Testing only
        raise RefreshFailedException('Could not refresh clamav db. returned exit code {}. Output = {}'.format(status, output))


def parse_refresh_output(output: str) -> dict:
    """
    Parse the refresh output to return a dict of db versions post-refresh
    :param output:
    :return:
    """

    # Example output:
    # Database test passed.
    # daily.cvd updated (version: 25889, sigs: 3716328, f - level: 63, builder: raynman)
    # main.cvd database is up to date (version: 59, sigs: 4564902, f - level: 60, builder: sigmgr)
    # bytecode.cvd database is up to date (version: 331, sigs: 94, f - level: 63, builder: anvilleg)

    daily_re = re.compile(r"\W*daily\.c[lv]d .*\(version: ([0-9]+),.*")
    main_re = re.compile(r"\W*main\.c[lv]d .*\(version: ([0-9]+),.*")
    bytecode_re = re.compile(r"\W*bytecode\.c[lv]d .*\(version: ([0-9]+),.*")

    daily_version = ""
    main_version = ""
    bytecode_version = ""

    for line in output.splitlines():
        m = daily_re.match(line)
        if m:
            daily_version = m.group(1)
            continue

        m = main_re.match(line)
        if m:
            main_version = m.group(1)
            continue

        m = bytecode_re.match(line)
        if m:
            bytecode_version = m.group(1)

    return {
        'daily': daily_version,
        'main': main_version,
        'bytecode': bytecode_version
    }


def parse_clamscan_output_line(scanned_file: str, line: str) -> tuple:
    """
    Parse a single line from the output

    :param scanned_file: full path of tarball scanned, for context
    :param line:
    :return:
    """
    regex = r'{}\!POSIX_TAR:.+\.\.\.\!\([0-9]\)POSIX_TAR:(.+):\s+(\S+)\s+FOUND$'.format(re.escape(scanned_file))
    line_match = re.match(regex, line)

    if line_match is not None:
        # Prefix the leading slash since that is omitted due to inspecting a tar, but the tar is a rootfs for an image
        path = "/" + line_match.groups()[-2]
        signature = line_match.groups()[-1]

        return path, signature

    return tuple()


def parse_clamscan(scanned_file: str, raw_output: str) -> list:
    """
    Split lines into k/v pairs.

    Example output:
    /analysis_scratch/a5346cac-802a-4cb9-a0c5-756b9b51858b/squashed.tar: Unix.Trojan.MSShellcode-40 FOUND
    /analysis_scratch/a5346cac-802a-4cb9-a0c5-756b9b51858b/squashed.tar!POSIX_TAR:bin/busybox!...!(8)POSIX_TAR:elf_payload1: Unix.Trojan.MSShellcode-40 FOUND"

    Each line is formatted as: `<path>: <signature> FOUND`

    With the --archive-verbose flag on clamscan, it will output the path within the archive that the item is found

    # returns list of dicts mapping path to signature e.g.
    # {
    # "/file": ["signature1", "signature2"]
    # }

    returns a list of objects [ {"path": "/somepath", "finding": <some signature for scanner>} ]

    :param scanned_file: str full path of file that was scanned
    :param raw_output:
    :return: list of dicts
    """

    summary_regex = re.compile('-+ SUMMARY -+')
    results = []

    for line in raw_output.splitlines():
        if summary_regex.match(line):
            # Output complete
            break

        result = parse_clamscan_output_line(scanned_file, line)
        if result:
            # Prefix each with "/" since it came from a tar file that is the rootfs
            results.append({"path": result[0], "signature": result[1]})
    return results


class ScanResult:
    def __init__(self, name: str, metadata: dict, findings: list):
        self.findings = findings
        self.name = name
        self.metadata = metadata

    def to_json(self):
        return {
            "scanner": self.name,
            "findings": self.findings,
            "metadata": self.metadata
        }


class MalwareScanner:
    """
    Generic class for implementing malware scanners that can be easily run the analyzer
    """
    name = None

    def __init__(self, config, squashed_tar_path, tempdir):
        """

        :param config: the config dict for this named scanner
        :param squashed_tar_path: the full file path to the squash.tar to be scanned
        :param tempdir: path to use as temp space during scans
        """
        self.config = config
        self.tar_path = squashed_tar_path
        self.enabled = self.config.get(self.name, {}).get('enabled', False) if self.config is not None and self.config.get(self.name) is not None and self.tar_path is not None else False
        self.refresh_enabled = self.config.get(self.name, {}).get('db_update_enabled', True) if self.config is not None and self.config.get(self.name) is not None else True
        self.tempdir = tempdir

    def run(self) -> ScanResult:
        raise NotImplementedError


class ClamAVRunner(MalwareScanner):
    name = 'clamav'

    def scan(self) -> list:
        """
        Run the scan, returning a list of findings

        :return:
        """
        tmpdir = self.tempdir
        output = run_clamav(self.tar_path, tempdir=tmpdir, dbdir=dbdir)
        print("Raw ClamAV stdout: " + output)
        return parse_clamscan(self.tar_path, output)

    def refresh_db(self) -> dict:
        """
        Update the signature db, and return the versions
        :return: version dict
        """
        refresh_output = refresh_clamavdb(databasedir=dbdir, configfile=config_path)
        return parse_refresh_output(refresh_output)


    def run(self):
        print("running clamav scan")

        if self.enabled:
            # Run clamav scan
            if self.refresh_enabled:
                db_versions = self.refresh_db()
            else:
                db_versions = {}

            meta = {'db_version': db_versions, 'db_update_enabled': self.refresh_enabled}
            print('ClamAV scan metadata: {}'.format(meta))

            output = self.scan()
            print('ClamAV findings: {}'.format(output))
        else:
            output = None
            meta = None

        return ScanResult(name=self.name, metadata=meta, findings=output)
