import os
import json
import re
import base64
import subprocess
from elem.exploit import ExploitSource
from elem.exploit import CuratedExploit

class CurationManager(object):
    """
    The primary interface to the curation information.

    Attributes:
        curation_location (str): The location of the curation information.
    """
    def __init__(self, curation_dir):
        self.curation_dir = curation_dir
        self.sources = dict()
        if not os.path.isdir(self.curation_dir):
            raise OSError("Curation directory not found: {0}".format(self.curation_dir))

    def add_source(self, source_name, location):
        if source_name not in self.sources.keys():
            self.modify_source(source_name, location)

    def delete_source(self, source_name):
        self.sources.pop(source_name)

    def modify_source(self, source_name, location):
        self.sources[source_name] = ExploitSource(source_name, location)

    def update_exploits(self, source_name='', all_exploits=False):
        exploits = []
        if source_name:
            exploits = CurationManager.load_exploits_from_source(self.sources[source_name],
                                                                 self.curation_dir,
                                                                 all_exploits)
        else:
            for _, source in self.sources.iteritems():
                exploits = CurationManager.load_exploits_from_source(source,
                                                                     self.curation_dir,
                                                                     all_exploits)
        for exploit in exploits:
            curated_exploit = CuratedExploit.from_eid(self.curation_dir, exploit['source'], exploit['id'])
            curated_exploit.reconcile(exploit)
            curated_exploit.write()

    def as_list(self):
        data = []
        filenames = [os.path.join(d, x)
                     for d, _, files in os.walk(self.curation_dir)
                     for x in files]
        for filename in filenames:
            with open(filename, 'r') as file_obj:
                try:
                    json_from_file = json.load(file_obj)
                    data.append(json_from_file)
                except ValueError:
                    pass
        return data

    def score(self, source, eid, cpe, kind, value):
        curated_exploit = CuratedExploit.from_eid(self.curation_dir, source, eid)
        curated_exploit.score(cpe, kind, value)
        curated_exploit.write()

    def set_stage(self, source, eid, cpe, command=None, packages=None, services=None, selinux=None, filename=None):
        curated_exploit = CuratedExploit.from_eid(self.curation_dir, source, eid)

        if command:
            curated_exploit.set_command(cpe, command)

        if packages:
            curated_exploit.set_packages(cpe, packages)

        if services:
            curated_exploit.set_services(cpe, services)

        if selinux:
            curated_exploit.set_selinux(cpe, selinux)

        if filename:
            curated_exploit.set_filename(cpe, os.path.basename(filename))

        curated_exploit.write()

    def copy(self, source, eid, destination, cpe, stage=False):
        curated_exploit = CuratedExploit.from_eid(self.curation_dir, source, eid)
        
        if os.path.isdir(destination):
            destination = os.path.join(destination, curated_exploit.get_filename(cpe))
        with open(destination, 'w') as file_obj:
            file_obj.write(base64.b64decode(curated_exploit.get_exploit()))

        if stage:
            self.stage(curated_exploit, cpe, os.path.dirname(destination))

    def stage(self, curated_exploit, cpe, destination_folder):

        if isinstance(curated_exploit.get_command(cpe), list):
            command = ' '.join(curated_exploit.get_command())
        else:
            command = curated_exploit.get_command(cpe)

        p = subprocess.Popen(command,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE,
                             cwd=destination_folder,
                             shell=True)
        out, err = p.communicate()
        lines = out.split('\n')
        error_lines = err.split('\n')
        if p.returncode != 0:
            raise OSError(error_lines)

    def cves_from_exploits(self, source, eids):
        cves = []
        for eid in eids:
            curated_exploit = CuratedExploit.from_eid(self.curation_dir, source, eid)
            cves = cves + curated_exploit.cves()

        return cves

    def csv(self, source='', cves=None, cpes=None, score_kind='', score_regex='', eid=''):

        lines = []
        data = self.as_list()
        data = self._filter(data, source, cves, cpes, score_kind, score_regex, eid)

        for curation_info in data:
            for cve in curation_info['cves']:

                if 'cpes' in curation_info.keys():
                    for cpe in curation_info['cpes'].keys():
                        for score in curation_info['cpes'][cpe]['scores'].keys():
                            line = [curation_info['source']]
                            line.append(curation_info['id'])
                            line.append(cve)
                            line.append(cpe)
                            line.append(score)
                            line.append(curation_info['cpes'][cpe]['scores'][score])
                            lines.append(','.join(line))
                else:
                    line = [curation_info['source']]
                    line.append(curation_info['id'])
                    lines.append(','.join(line))

        return '\n'.join(lines)

    @classmethod
    def _filter(cls, data, source='', cves=None, cpes=None, score_kind='', score_regex='', eid=''):

        if cves:
            data = [curation_info for curation_info in data if CurationManager._cves_in_exploit(curation_info, cves)]

        if cpes:
            data = [curation_info for curation_info in data if CurationManager._cpes_in_exploit(curation_info, cpes)]

        if source:
            data = [curation_info for curation_info in data if source == curation_info['source']]

        if eid:
            data = [curation_info for curation_info in data if eid == curation_info['id']]

        if score_kind:
            data = [curation_info for curation_info in data if CurationManager._score_kind_in_exploit(curation_info, score_kind)]

        if score_regex:
            data = [curation_info for curation_info in data if CurationManager._score_regex_in_exploit(curation_info, score_regex)]

        return data

    @classmethod
    def load_exploits_from_source(cls, exploit_source, curation_dir, all_exploits=False):
        if all_exploits:
            exploits = [exploit for exploit in exploit_source.exploits()
                        if exploit['cves']]
        else:
            exploits = [exploit for exploit in exploit_source.exploits()
                        if exploit['cves']
                        if not os.path.isfile(os.path.join(curation_dir,
                                                           exploit['source'],
                                                           exploit['id'] + '.json'))]
        return exploits

    @classmethod
    def _cves_in_exploit(cls, exploit, cves):
        for cve in cves:
            if cve in exploit['cves']:
                return True
        return False

    @classmethod
    def _cpes_in_exploit(cls, exploit, cpes):
        if 'cpes' in exploit.keys():
            for cpe in cpes:
                if cpe in exploit['cpes'].keys():
                    return True
        return False

    @classmethod
    def _score_kind_in_exploit(cls, exploit, kind):
        if 'cpes' in exploit.keys():
            for cpe in exploit['cpes'].keys():
                if kind in exploit['cpes'][cpe]['scores'].keys():
                    return True
        return False

    @classmethod
    def _score_regex_in_exploit(cls, exploit, regex):
        if 'cpes' in exploit.keys():
            for cpe in exploit['cpes'].keys():
                for score in exploit['cpes'][cpe]['scores'].keys():
                    if re.match(regex, exploit['cpes'][cpe]['scores'][score]):
                        return True
        return False
