#!/usr/bin/env python
# Copyright (c) 2021 VMware, Inc. All Rights Reserved.
# SPDX-License-Identifier: BSD-2 License
# The full license information can be found in LICENSE.txt
# in the root directory of this project.

import collections
import logging
import os
import re
import uuid

from lydian.apps.base import exposify
from lydian.apps.console import Console
from lydian.apps.internal.setup import SetupDB


log = logging.getLogger(__name__)


@exposify
class Rapidscan(Console):
    NAME = "Rapidscan"
    RUNDIR = '/home/kali'
    TOOL_PATH = '/home/kali'
    TOOL_BINARY = 'rapidscan.py'
    CMD_PREFIX = 'su kali -c'

    def __init__(self):
        super().__init__()
        self.requests = {}
        self.targets = collections.defaultdict()

    def start_run(self, *args, reqid=None, **kwargs):
        """
        """
        log.info("args=%s kwargs=%s", args, kwargs)
        if reqid in self.requests:
            log.info('%s - scan for %s is already running', self.__class__.__name__, reqid)
            return
        reqid = reqid or '%s' % uuid.uuid4()
        log.info("Running in directory %s", os.getcwd())
        cmnd = '%s "%s/%s ' % (self.CMD_PREFIX, self.TOOL_PATH, self.TOOL_BINARY)
        cmnd += ' '.join(args) + ' ' if args else ' '
        cmnd += ' '.join("--{}={}".format(kw, val) for kw, val in kwargs.items())
        cmnd += '"'
        try:
            p = self._start_subprocess(cmnd, cwd=self.RUNDIR)
        except Exception as e:
            log.error(e)
        log.info("Running command %s: %s", cmnd, p.pid)
        self.requests[reqid] = p
        self.targets[reqid] = args[0]
        return reqid

    def is_running(self):
        """
        Tells if app is available.
        """
        return True

    def show_running_processes(self):
        if not self.requests:
            return None
        return {req: p.pid for req, p in self.requests.items()}

    def is_run_complete(self, reqid):
        log.info("Checking run for remote %s", reqid)
        if reqid not in self.requests:
            log.info('% - Invalid request %s', self.__class__.__name__, reqid)
            return None
        return not self._is_alive(self.requests[reqid])

    def stop_run(self, reqid):
        if reqid not in self.requests:
            log.info('%s - Invalid request %s', self.__class__.__name__, reqid)
            return None
        if self._is_alive(self.requests[reqid]): # Still running
            self._kill_subprocess(self.requests[reqid], True)

    def get_result(self, reqid, **kwargs):
        """
        This method should be implemented by child class.
        The DB is set to a separate table for each application./
        """
        with SetupDB() as db:
            db.table = db.RAPIDSCAN
            return db.read(reqid=reqid, **kwargs)

    def save_results_to_db(self, reqid):
        proc = self.requests[reqid]
        if not proc:
            return
        stdout_val = proc.communicate()[0]
        output = stdout_val.strip().decode('utf-8')
        target = self.targets[reqid]
        assert target, "No target available for reqid %s" % reqid
        filtered_data = self._parse(output.split("\n"))
        log.info(filtered_data)
        with SetupDB() as db:
            db.table = db.RAPIDSCAN
            for tool, info in filtered_data.items():
                db.write(
                    host=target,
                    tool=tool,
                    reqid=reqid,
                    severity=info['severity'],
                    message=info['message'])

    def _parse(self, results):
        SEARCH = 0
        FOUND = 1
        state = SEARCH
        vul_info = {}
        prev_line = results[0]
        tool = severity = ""
        description = ""
        for line in results[1:]:
            if line == "":
                continue
            line = re.sub(r'\x1b\[[0-9]{1,2}m', '', line)
            line = re.sub(r'[\\/|]*\x08', '', line)
            line = re.sub(r'\x1b\[41m[\\/|]', '', line)
            line = re.sub(r'\x1b\[0m\x08', '', line)
            line = re.sub(r'\x1b[\[FK0-9]+', '', line)
            if "Vulnerability Threat Level" in line and state == SEARCH:
                state = FOUND
                log.debug("Line: [%s]\nPrev Line: [%s]", line, prev_line)
                tool = prev_line.split('|')[1].split()[0].strip()
            elif state == FOUND:
                log.debug("[%s]", line)
                severity, description = line.strip().split(' ', 1)
                state = SEARCH
            elif state != SEARCH:
                log.warning("Unexpected pattern. Please check the logs for formatting\n%s", line)
            prev_line = line
            if severity and tool:
                vul_info[tool] = {"severity": severity, "message": description}
        return vul_info
