#!/usr/bin/env python
#
# Author: Fredrik Hedman <fredrik.hedman@noruna.se>
# VERSION: 1.20.1
# LICENSE: MIT LICENSE
#
"""The VulnerabilityVector class encapsulates a vulnerability vector.

This concept is described in the [CVSS
Guide.](http://www.first.org/cvss/cvss-guide.html)

A VulnerabilityVector is initialized with a string containing a
vulnerability vector.  The operations help in checking, modifying and
using the vector.

A vector is a list of components separated by a slash ("/").  Each
component is is a pair of values separated by a colon (":"). The first
value is the abbreviated metric name and the second value is the
abbreviated metric value.

The base, temporal, and environmental vectors are shown below with the
possible metric values listed within square brackets ([...])

* Base
"AV:[L,A,N]/AC:[H,M,L]/Au:[M,S,N]/C:[N,P,C]/I:[N,P,C]/A:[N,P,C]"

* Temporal
"E:[U,POC,F,H,ND]/RL:[OF,TF,W,U,ND]/RC:[UC,UR,C,ND]"

* Environmental
"CDP:[N,L,LM,MH,H,ND]/TD:[N,L,M,H,ND]/CR:[L,M,H,ND]/"
"IR:[L,M,H,ND]/AR:[L,M,H,ND]"

The vector lists these metrics in a predetermined order.

When a temporal or environmental metric is not used it is assigned the
value "ND" (not defined).

"""
from collections import OrderedDict
from .metric import Metric


def parse_vector_template(txt):
    """Helper to parse specification."""
    keys = [k.split(':')[0] for k in txt.split('/')]
    vals = [k.split(':')[1][1:-1].split(',') for k in txt.split('/')]
    return OrderedDict(zip(keys, vals))

BASE_SCORE = "AV:[L,A,N]/AC:[H,M,L]/Au:[M,S,N]/C:[N,P,C]/I:[N,P,C]/A:[N,P,C]"
TEMPORAL_SCORE = "E:[U,POC,F,H,ND]/RL:[OF,TF,W,U,ND]/RC:[UC,UR,C,ND]"
ENVIRONMENTAL_SCORE = ("CDP:[N,L,LM,MH,H,ND]/TD:[N,L,M,H,ND]/CR:[L,M,H,ND]/"
                       "IR:[L,M,H,ND]/AR:[L,M,H,ND]")

BASE = parse_vector_template(BASE_SCORE)
TEMPORAL = parse_vector_template(TEMPORAL_SCORE)
ENVIRONMENTAL = parse_vector_template(ENVIRONMENTAL_SCORE)


def all_metrics():
    """Build and return a list of all metrics."""
    ll = []
    ll.extend(base_metrics())
    ll.extend(temporal_metrics())
    ll.extend(environmental_metrics())
    return ll


def base_metrics():
    """Wrap base metrics data and possible values."""
    metrics = [
        ["Access Vector", "AV",
         [("Local", "L", 0.395, "Local access"),
          ("Adjecent Network", "A", 0.646, "Adjacent network access"),
          ("Network", "N", 1.0, "Network access"), ]],
        ["Access Complexity", "AC",
         [("High", "H", 0.35, "Specialized access conditions exist"),
          ("Medium", "M", 0.61,
           "The access conditions are somewhat specialized"),
          ("Low", "L", 0.71, "No specialized access exist"), ]],
        ["Authentication", "Au",
         [("None", "N", 0.704, "Authentication not required"),
          ("Multiple", "M", 0.45, "Authenticate two or more times"),
          ("Single", "S", 0.56, "Logged into the system"), ]],
        ["Confidentiality Impact", "C",
         [("None", "N", 0.0, "No impact"),
          ("Partial", "P", 0.275, "Considerable disclosure"),
          ("Complete", "C", 0.660, "Total inforamtion disclosure"), ]],
        ["Integrity Impact", "I",
         [("None", "N", 0.0, "No impact"),
          ("Partial", "P", 0.275,
           "Possible to modify some system files or information"),
          ("Complete", "C", 0.660, "Total compromise of system integrity"), ]],
        ["Availability Impact", "A",
         [("None", "N", 0.0, "No impact"),
          ("Partial", "P", 0.275,
           "Reduced performance or interruptions in resource availability"),
          ("Complete", "C", 0.660,
           "Total shutdown of the affected resource"), ]],
    ]
    return metrics


def temporal_metrics():
    """Wrap temporal metrics data and possible values."""
    metrics = [
        ["Exploitability", "E",
         [("Not Defined", "ND", 1.0, "Skip this metric"),
          ("Unproven", "U", 0.85, "No exploit code is available"),
          ("Proof-of-Concept", "POC", 0.9,
           "Proof-of-concept exploit code exists"),
          ("Functional", "F", 0.95, "Functional exploit code is available"),
          ("High", "H", 1.0,
           "Exploitable by functional mobile autonomous code"), ]],
        ["Remediation Level", "RL",
         [("Not Defined", "ND", 1.0,
           "Skip this metric"),
          ("Official Fix", "OF", 0.87,
           "Complete vendor solution is available"),
          ("Temporary Fix", "TF", 0.90,
           "Official but temporary fix available"),
          ("Workaround", "W", 0.95,
           "Unofficial, non-vendor solution available"),
          ("Unavailable", "U", 1.0,
           "No solution available or it is impossible to apply"), ]],
        ["Report Confidence", "RC",
         [("Not Defined", "ND", 1.0,
           "Skip this metric"),
          ("Unconfirmed", "UC", 0.90,
           "Single unconfirmed source"),
          ("Uncorroborated", "UR", 0.95,
           "Multiple non-official sources"),
          ("Confirmed", "C", 1.0,
           "Acknowledged by the vendor or author"), ]],
    ]
    return metrics


def environmental_metrics():
    """Wrap environmental metrics data and possible values."""
    metrics = [
        ["Collateral Damage Potential", "CDP",
         [("Not Defined", "ND", 0.0, "Skip this metric"),
          ("None", "N", 0.0, "No potential for loss of life"),
          ("Low", "L", 0.1,
           "Potential for slight physical or property damage"),
          ("Low-Medium", "LM", 0.3, "Moderate physical or property damage"),
          ("Medium-High", "MH", 0.4,
           "Significant physical or property damage or loss"),
          ("High", "H", 0.5,
           "Catastrophic physical or property damage and loss"), ]],
        ["Target Distribution", "TD",
         [("Not Defined", "ND", 1.0, "Skip this metric"),
          ("None", "N", 0.0, "No target systems exist"),
          ("Low", "L", 0.25,
           "Targets exist on a small scale inside the environment"),
          ("Medium", "M", 0.75, "Targets exist on a medium scale"),
          ("High", "H", 1.0, "Targets exist on a considerable scale"), ]],
        ["Confidentiality Requirement", "CR",
         [("Not Defined", "ND", 1.0, "Skip this metric"),
          ("Low", "L", 0.5, "Limited adverse effect"),
          ("Medium", "M", 1.0, "Serious adverse effect"),
          ("High", "H", 1.51, "Catastrophic adverse effect"), ]],
        ["Integrity Requirement", "IR",
         [("Not Defined", "ND", 1.0, "Skip this metric"),
          ("Low", "L", 0.5, "Limited adverse effect"),
          ("Medium", "M", 1.0, "Serious adverse effect"),
          ("High", "H", 1.51, "Catastrophic adverse effect"), ]],
        ["Availability Requirement", "AR",
         [("Not Defined", "ND", 1.0, "Skip this metric"),
          ("Low", "L", 0.5, "Limited adverse effect"),
          ("Medium", "M", 1.0, "Serious adverse effect"),
          ("High", "H", 1.51, "Catastrophic adverse effect"), ]],
    ]
    return metrics


def add_padding(to_length, selected):
    """Extend selected to_length elements."""
    if selected is None:
        selected = []
    padding = to_length - len(selected)
    if padding > 0:
        selected.extend(padding * [None])
    return selected


def prepare_metrics(L, selected):
    """Prepare a list of selected metrics."""
    lmetrics = []
    for ii, mm in enumerate(L):
        lmetrics.append(Metric(*mm, index=selected[ii]))
    return lmetrics


def cvs_factory(cls, selected=None):
    """Common Vulnerability Score factory."""
    ll = all_metrics()
    selected = add_padding(len(ll), selected)
    lmetrics = prepare_metrics(ll, selected)
    return cls(lmetrics)


class InvalidBaseVector(Exception):
    pass


class VulnerabilityVectorError(Exception):
    pass


class VulnerabilityVector(object):

    """The VulnerabilityVector class encapsulates a vulnerability vector."""

    def __init__(self, seq):
        """Start from a vector of metrics."""
        self.vvec = self._parse_vulnerability_string(seq)
        self.vulnerability_string = seq
        self.base = BASE
        self.temporal = TEMPORAL
        self.environmental = ENVIRONMENTAL
        self.all = OrderedDict(BASE)
        self.all.update(TEMPORAL)
        self.all.update(ENVIRONMENTAL)

    def _parse_vulnerability_string(self, seq):
        try:
            d = OrderedDict(map(lambda x: x.split(":"), seq.split("/")))
            l = list(map(lambda x: x.split(":"), seq.split("/")))
            if len(d) != len(l):
                msg = "{0} duplicate elements" \
                      " in vector".format(seq)
                raise InvalidBaseVector(msg)
        except (InvalidBaseVector) as e:
            print('Error: invalid vulnerability vector.')
            print('Hint: {}'.format(e))
            raise
        except (ValueError) as e:
            msg00 = "Error in parsing vulnerability vector {0}".format(seq)
            print('{1} : {0}'.format(e, InvalidBaseVector(msg00)))
            raise
        return d

    def valid(self):
        """Validate vulnerability string.

        Each element is a key-value pair where the key is a Metric and
        the value is a MetricValue.  Check that each pair is valid.  The
        key should be one in the set of valid Metrics and the
        corresponding value should be in the set of valid MetricValues.

        """
        for key, value in self.vvec.items():
            try:
                if (key in self.all.keys()) and (value in self.all[key]):
                    continue
                elif key not in self.all.keys():
                    msg00 = "{0} not a metric".format(key)
                elif value not in self.all[key]:
                    msg01 = "Must be one of: {0})".format(self.all[key])
                    msg00 = "{0} not a metric value. {1}".format(value, msg01)
                raise InvalidBaseVector(msg00)
            except InvalidBaseVector as e:
                print('Error: invalid vulnerability vector.')
                print('Hint: {0}'.format(e))
                raise
        return self

    def complete(self):
        """Ensure vulnerability string is complete.

        In addition to being a valid vulnerability string, the number
        of elements and the order of elements should correspond to the
        expected order represented by the metric_seq.

        """
        self.valid()
        try:
            self._valid_length(self.vvec)
            self._valid_order(self.vvec)
        except (InvalidBaseVector) as e:
            print('Error: invalid vulnerability vector.')
            print('Hint: {}'.format(e))
            raise
        return self

    def _valid_length(self, vec):
        valid_base_len = len(self.base)
        valid_temp_len = valid_base_len + len(self.temporal)
        valid_env_len = valid_temp_len + len(self.environmental)
        # print("##", len(vec), valid_base_len, valid_temp_len, valid_env_len)
        if (valid_base_len != len(vec) and
                valid_temp_len != len(vec) and valid_env_len != len(vec)):
                    msg = "wrong number of elements in vector"
                    raise InvalidBaseVector(msg)
        return True

    def _valid_order(self, vec):
        base_metric_index_order = list(self.all.keys())
        for ii, v in enumerate(vec):
            if v != base_metric_index_order[ii]:
                msg = "{0} duplicate elements or" \
                      " incorrect ordering in vector".format(vec)
                raise InvalidBaseVector(msg)
        return True

    def metric_values(self):
        selected = []
        for idx, value in self.vvec.items():
            selected.append(value)
        return selected


if __name__ == "__main__":
    import doctest
    doctest.testmod()
