from report_ranger import validation
from report_ranger.utils.jinja_helpers import log_jinja2_error
from report_ranger.utils.mdread import markdown_from_file, process_template
from report_ranger.output_formatter.outputformatter import oxfordcomma
import os
import logging
import copy
import datetime
from collections import Counter
from num2words import num2words
import jinja2
import traceback
import cerberus
from report_ranger.helpers import make_list

log = logging.getLogger(__name__)


def vuln_from_file(file_loc, riskassessment, env=None, vulnerability_validation={}, scope=None, watcher=None):
    ''' Get the vulnerability markdown from a file, process the headers and make sure they're there. Return a Vulnerability object.'''
    log.debug(f"Loading vulnerability file {file_loc}")
    try:
        headers, vulnmarkdown = markdown_from_file(file_loc, env=env, watcher=watcher)
    except:
        log.warn(f"Could not read vulnerability file {file_loc}")

    if 'name' in headers:
        if not isinstance(headers['name'], str):
            log.warn(f"Vulnerability name not a string for file {file_loc}")
            name = file_loc
            headers['name'] = file_loc
        else:
            name = headers['name']
    elif 'Name' in headers:  # This is for legacy reasons, it used to be case insensitive so we still accept "Name"
        name = headers['Name']
        headers['name'] = headers['Name']
    else:
        log.warn("Vulnerability name not in headers for file {}".format(file_loc))
        name = file_loc
        headers['name'] = file_loc

    if scope != None:
        headers['scope'] = scope

    validation.validate_headers(
        vulnerability_validation, headers, validation.default_vulnerability_validation, name=name)

    # Generate the new env variables they can play with
    headers['cwd'] = os.path.dirname(os.path.join(
        os.path.curdir, file_loc))  # for including screenshots

    return Vulnerability(name, headers=headers, markdown=vulnmarkdown, riskassessment=riskassessment, filename=file_loc)


class VulnerabilityList:
    """Reads a directory for vulnerability definition files, arranges them, and stores the vulnerabilities as objects.

    The class is usually referenced inside the markdown file itself to display the vulnerability information. It keeps the following attributes:

    - scopes : a List of Scope objects corresponding to the vulnerability scopes
    """

    def __init__(self):
        self.scopes = []

    def add_from_dir(self, vulndir, env, riskassessment, vulnerability_validation={}, watcher=None):
        # Make sure it's a directory
        if not os.path.isdir(vulndir):
            raise FileNotFoundError("Directory '{0}' doesn't exist. Skipping import of vulnerabilities.".format(
                vulndir))

        log.debug("Reading vulnerabilities")
        vulnerabilities = []

        scopeintroductions = {}
        log.debug(f"Vulnerability add_from_dir WATCHER: {watcher}")

        # Get the vulnerability writeups from the vulndir
        for f in os.listdir(vulndir):
            if watcher:
                watcher.add_path(vulndir)
            # Scopes can be split into directories within the vulndir
            if os.path.isdir(os.path.join(vulndir, f)):
                if watcher:
                    watcher.add_path(os.path.join(vulndir, f))
                log.debug(
                    "* New vulnerability scope in directory '{}'".format((os.path.join(vulndir, f))))
                vdir = os.path.join(vulndir, f)
                for v in os.listdir(vdir):
                    if v == "template.md" or v == "template.md.rr":
                        continue
                    if v == "introduction.md" or v == "introduction.md.rr":
                        try:
                            scopeintroductions[f] = markdown_from_file(
                                os.path.join(vdir, v), env=env, watcher=watcher)
                        except:
                            log.warn(f"Could not read introduction file {v} in vulnerability directory {vdir}")
                        continue
                    fn, ext = os.path.splitext(v)
                    if ext == ".md" or ext == ".rr":  # Only process markdown files
                        filename = os.path.join(vdir, v)
                        log.debug(
                            "  - Reading vulnerability file '{}'".format(filename))
                        try:
                            vuln = vuln_from_file(
                                filename, riskassessment=riskassessment, env=env, vulnerability_validation=vulnerability_validation, scope=f, watcher=watcher)
                            vulnerabilities += [vuln]
                        except validation.ValidationError as e:
                            pass  # ValidationError already raises logs
                        except Exception as e:
                            log.error(
                                "Received exception when processing markdown for vulnerability {}: {}".format(filename, e.args))
                            log.error(
                                "Removing markdown for vulnerability {}".format(self.name))
                            self.markdown = ""
            else:  # It's a file
                if f == "template.md" or f == "template.md.rr":
                    continue
                if f == "introduction.md" or f == "introduction.md.rr":
                    try:
                        scopeintroduction = markdown_from_file(
                            os.path.join(vulndir, f), env=env, watcher=watcher)
                        if "scope" in scopeintroduction[0]:
                            # There's a scope defined in the headers of the introduction
                            scopeintroductions[f] = scopeintroduction
                        else:
                            scopeintroductions[''] = scopeintroduction
                    except:
                        log.warn(f"Could not read introduction file {v} in vulnerability directory {vdir}")
                    continue
                    
                log.debug(
                    "* Reading vulnerability file '{}'".format(os.path.join(vulndir, f)))
                fn, ext = os.path.splitext(f)
                if ext == ".md" or ext == ".rr":  # Only process markdown files
                    try:
                        vulnerabilities += [vuln_from_file(
                            os.path.join(vulndir, f), riskassessment=riskassessment, env=env, watcher=watcher)]
                    except validation.ValidationError as e:
                        pass  # ValidationError already raises logs
                    except Exception as e:
                        log.error(
                            "Received exception when processing markdown for vulnerability {}: {}".format(filename, e.args))
                        log.error(
                            "Removing markdown for vulnerability {}".format(self.name))
                        self.markdown = ""

        # Sort the vulnerabilities into scopes
        scopes = {}

        # This is a marker for the highest risk
        highestrisk = riskassessment.risks[0]
        nhighestrisk = 0

        suggest_rc = set()
        suggest_ar = set()

        # Suggest high level findings
        for vuln in vulnerabilities:
            if 'suggest_rc' in vuln.headers:
                suggest_rc.update(vuln.headers['suggest_rc'])
            if 'suggest_ar' in vuln.headers:
                suggest_ar.update(vuln.headers['suggest_ar'])

        for rc in suggest_rc:
            log.info("Suggested root cause: {}".format(rc))
        for ar in suggest_ar:
            log.info("Suggested additional recommendation: {}".format(ar))

        for i in vulnerabilities:
            if i.nrisk > nhighestrisk:  # We have a new highest risk!
                nhighestrisk = i.nrisk
                highestrisk = i.risk
            if not i.scope in scopes:  # New scope
                scopes[i.scope] = Scope(i.scope)
            scopes[i.scope].add_vulnerability(i)

        for scope, introduction in scopeintroductions.items():
            if scope not in scopes:
                scopes[scope] = Scope(scope)
            scopes[scope].set_introduction(introduction[0], introduction[1])

        # Add the highest risk to the variables
        self.highestrisk = highestrisk
        env.set_variable('highestrisk', highestrisk)
        log.debug("Highest risk of any vuln is {}".format(highestrisk))

        self.scopes = list(scopes.values())
        self.scopes.sort()

        # Assign references
        ref = 1
        for scope in self.scopes:
            for vulnerability in scope.vulnerabilities:
                vulnerability.ref = ref
                ref += 1

    def get_scope_list(self, validation=None, date=None):
        scopelist = {}
        if date:
            updatedvlist = self.updated(date)
            scopes = updatedvlist.scopes
        else:
            scopes = self.scopes
        for scope in scopes:
            scopelist[scope.name] = scope.get_list(validation)
        return scopelist

    def get_ref(self, vulnerability_name):
        """ Get the reference of a vulnerability by name. Usually called from a markdown file. """
        for scope in self.scopes:
            ref = scope.get_ref(vulnerability_name)
            if ref:
                return ref
        return ""

    def updated(self, date=datetime.date.today()):
        # Alias for string today
        if date == "today" or date == "now":
            date = datetime.date.today()

        newlist = VulnerabilityList()
        for scope in self.scopes:
            newlist.scopes.append(scope.updated(date))
        newlist.scopes.sort()
        return newlist

    def nvulns(self, riskrating=""):
        """ Get the number of vulnerabilities with the given risk rating """
        return sum(scope.nvulns(riskrating) for scope in self.scopes)

    def nriskdict(self):
        """ Get a dictionary with risk:nrisks for all vulnerabilities """
        nvulnlist = {}
        for scope in self.scopes:
            nvulnlist.update(scope.nriskdict())
        return nvulnlist

    def nrisklist(self):
        return oxfordcomma(["{} ({}) {} risk".format(num2words(nrisk), nrisk, risk)
                            for risk, nrisk in self.nriskdict().items()])

    def generate_markdown(self, env):
        for scope in self.scopes:
            scope.generate_markdown(env)


class Scope:
    """ Holds a vulnerability scope, a container for vulnerabilities.

    Has the following attributes:

    name : The name of the scope as a string
    vulnerabilities : A List of Vulnerability objects, corresponding to the vulnerabilities in this scope.
                      Should always be sorted by severity and weight.
    highestrisk : The highest risk in this scope as a string
    sectionweight : The weight as a number used for ordering the scopes in a VulnerabilityList (highest weight first).
                    Called sectionweight to have consistency with sections.
    introduction_headers : The environment headers of an introduction, used for Jinja templating the markdown
    introduction : The introduction markdown
    """

    def __init__(self, name, sectionweight=0):
        self.name = name
        self.vulnerabilities = []
        self.highestrisk = 0
        self.sectionweight = sectionweight
        self.introduction_headers = {}
        self.introduction = ""

        # Process the introduction section weight if it's there
        if "sectionweight" in self.introduction_headers and self.introduction_headers["sectionweight"] > self.sectionweight:
            log.info(f"Setting sectionweight for {name} to {self.introduction_headers['sectionweight']} based on introduction header")
            self.sectionweight = self.introduction_headers["sectionweight"]

    def __lt__(self, other):
        if other.sectionweight == self.sectionweight:
            return other.highestrisk < self.highestrisk
        else:
            return other.sectionweight < self.sectionweight

    def add_vulnerability(self, vulnerability):
        """ Add a vulnerability to the scope. Automatically sort the vulnerabilities after adding. """
        self.vulnerabilities.append(vulnerability)
        if vulnerability.nrisk > self.highestrisk:
            self.highestrisk = vulnerability.nrisk
        if vulnerability.sectionweight > self.sectionweight:
            self.sectionweight = vulnerability.sectionweight

        self.vulnerabilities.sort()
    
    def set_introduction(self, introduction_headers, introduction):
        self.introduction_headers = introduction_headers
        self.introduction = introduction
        # Process the introduction section weight if it's there
        if "sectionweight" in self.introduction_headers and self.introduction_headers["sectionweight"] > self.sectionweight:
            log.info(f"Setting sectionweight for {self.name} to {self.introduction_headers['sectionweight']} based on introduction header")
            self.sectionweight = self.introduction_headers["sectionweight"]

    def get_list(self, validation=None):
        vulnlist = []
        for vuln in self.vulnerabilities:
            if validation == None:
                vulnlist.append(vuln.get_list())
            else:
                vl = vuln.get_list()
                validator = cerberus.Validator()
                validator.allow_unknown = True
                if vuln.validate(vl):
                    vulnlist.append(vuln.get_list())
        return vulnlist

    def get_ref(self, vulnerability_name):
        """ Get the reference of a vulnerability by name. Usually called from a markdown file. """
        for vulnerability in self.vulnerabilities:
            if vulnerability_name.lower() in vulnerability.name.lower():
                return vulnerability.ref

        return None

    def nvulns(self, riskrating=""):
        """ Get the number of vulnerabilities with the given risk rating """
        return sum(riskrating in vuln.risk for vuln in self.vulnerabilities)

    def nriskdict(self):
        """ Get a dictionary with risk:nrisks for all vulnerabilities """
        return Counter(vulnerability.risk for vulnerability in self.vulnerabilities)

    def updated(self, date=datetime.date.today()):
        """ Process the 'updates' header """
        newscope = copy.copy(self)

        # Regenerate with the added vulnerabilities
        newscope.vulnerabilities = []
        for vuln in self.vulnerabilities:
            newscope.add_vulnerability(vuln.updated(date))

        return newscope

    def generate_markdown(self, env):
        for vuln in self.vulnerabilities:
            vuln.generate_markdown(env)
        if self.introduction != "":
            self.introduction = process_template(self.introduction_headers, self.introduction, env)


class Vulnerability:
    """A class for vulnerability markdown.

    Has the following attributes:

     - name : The name of the vulnerability as a string
     - vulnerability_markdown : The markdown of the vulnerability as a tuple (headers, markdown)
     - scope : The name of the scope as a string. Used by VulnerabilityList for sorting a bucket of vulnerabilities into scopes,
                 taken from the headers, but this may be overwritten if it's in a scope bucket already.
     - weight : A number, used for sorting vulnerabilities. Taken from the headers.
     - sectionweight : A number, used for increasing the weight of a scope for sorting them. Set for particularly bad vulns. Taken from the headers.
     - risk : The risk of the vuln as a string.
     - nrisk : The risk of the vuln as a number. Used for sorting vulnerabilities.
     - headers : The headers of the markdown.
    """

    def __init__(self, name, riskassessment, headers=dict(), markdown='', scope='', weight=0, sectionweight=0, filename=''):

        self.name = name
        self.vulnerability_markdown = markdown
        self.scope = scope
        self.ref = 0  # ref will be set by VulnerabilityList
        self.filename = filename
        self.riskassessment = riskassessment

        if "riskassessment" in headers:
            if riskassessment.id not in headers['riskassessment']:
                raise Exception("Risk assessment {} not in the risk assessment for vulnerability {}".format(
                    riskassessment.name, name))
            headers.update(headers['riskassessment'][riskassessment.id])

        if "scope" in headers:
            self.scope = headers["scope"]
        else:
            self.scope = ''

        if "weight" in headers:
            self.weight = float(headers["weight"])
        else:
            self.weight = 0

        if "sectionweight" in headers:
            self.sectionweight = float(headers["sectionweight"])
        else:
            self.sectionweight = 0

        self.risk = riskassessment.get_risk(headers)
        self.nrisk = next((i for i, x in enumerate(
            riskassessment.risks) if x == self.risk), 0)
        headers['risk'] = self.risk

        if "status" not in headers:
            headers['status'] = "Open"

        self.status = headers['status']
        self.originalstatus = self.status
        self.originalrisk = self.risk

        self.headers = headers

    def __lt__(self, other):
        if other.nrisk == self.nrisk:
            return other.weight < self.weight
        else:
            return other.nrisk < self.nrisk

    def get_list(self):
        vulnlist = {
            'name': self.name,
            'ref': self.ref,
            'status': self.status,
            'risk': self.risk,
            'nrisk': self.nrisk,
            'originalstatus': self.originalstatus,
            'originalrisk': self.originalrisk,
            'filename': self.filename,
            'headers': self.headers
        }
        for stage in self.riskassessment.stages:
            if stage['name'] in self.headers:
                vulnlist[stage['name']] = self.headers[stage['name']]
        return vulnlist

    def updated(self, date=datetime.date.today()):
        # Get a new vulnerability
        newvuln = copy.copy(self)

        # Deep headers needs a deep copy otherwise the updates will change the original vuln
        newvuln.headers = copy.deepcopy(self.headers)

        if 'updates' in self.headers:
            for update in make_list(self.headers['updates']):
                if 'date' not in update or update['date'] <= date or date == "all":
                    for newvar in update.keys():
                        if newvar == 'date':
                            continue
                        newvuln.headers[newvar] = update[newvar]

        # We need to re-evaluate the risk just in case the likelihood or consequence changed
        newvuln.risk = newvuln.riskassessment.get_risk(newvuln.headers)
        newvuln.headers['risk'] = newvuln.risk
        newvuln.status = newvuln.headers['status']

        return newvuln

    def generate_markdown(self, env):
        """ Get the markdown of the vulnerability, processing the template with the current env.

        Usually called from a markdown file. """
        # First update the ref
        self.headers['ref'] = self.ref
        self.headers['scope'] = self.scope

        try:
            self.markdown = process_template(self.headers, self.vulnerability_markdown,
                                             env, name="vulnerability {}".format(self.name), filename=self.filename)
        except jinja2.exceptions.TemplateSyntaxError as error:
            log.error("Jinja2 error processing vulnerability {}: {} at lineno {} for file {}".format(
                self.name, error.message, error.lineno, filename=self.filename))
            log.error("Removing markdown for vulnerability {}".format(self.name))
            log_jinja2_error(self.vulnerability_markdown, error)
            self.markdown = ""
        except Exception as e:
            log.error(
                "Received exception when processing markdown for vulnerability {}: {}".format(self.name, e.args))
            log.error(traceback.format_exc())
            log.error("Removing markdown for vulnerability {}".format(self.name))
            self.markdown = ""
