#!/usr/bin/env python3
  # -*- coding: utf-8 -*-
  """
  @author: chribell
  
  This script is used to create a simple CAPEC database as an SQLite file.
  
  Arguments
  ---------
  1. Directory path of CAPEC sqlite db
  2. Path of ouput sqlite file
  
  
  Example
  ---------
  ./create_capec_db.py /directory_of_xml output_capec_db
  
  """
  
  from sys import argv
  from os import path
  from sqlitedict import SqliteDict
  import xml.etree.ElementTree as ET
  
  
  def read_capec(capec_file):
      """
      Parameters
      ----------
      capec_file : str
          The capec.xml file path.
  
      Returns
      -------
      capecs : dict
          A CAPEC dictionary.
  
      """
      tree = ET.parse(capec_file)
      root = tree.getroot()
  
      namespace = '{http://capec.mitre.org/capec-3}'
      
      attack_patterns = root[0]
      capecs = {}
  
      for ap in attack_patterns:
          _id = ap.attrib['ID']
          _name = ap.attrib['Name']
          _description = ap.findall(namespace + 'Description')[0].text
          cwes = []
          relationships = {}
          for relation in ap.findall(namespace + 'Related_Attack_Patterns/' + namespace + 'Related_Attack_Pattern'):
              relationships[relation.attrib['CAPEC_ID']] = relation.attrib['Nature']
          for cwe in ap.findall(namespace + 'Related_Weaknesses/' + namespace + 'Related_Weakness'):
              cwes.append(cwe.attrib['CWE_ID'])
          parents = [k for k, v in relationships.items() if v == 'ChildOf']
          capecs[_id] = {
              'id': _id,
              'name': _name,
              'description':  _description,
              'cwe': cwes,
              'parents': parents,
              'relationships': relationships,
          }
      return capecs
  
  def read_cwe(cwe_file):
      """
      Parameters
      ----------
      cwe_file : str
          The cwe.xml file path.
  
      Returns
      -------
      cwes : dict
          A CWE dictionary.
  
      """
      tree = ET.parse(cwe_file)
      root = tree.getroot()
  
      namespace = '{http://cwe.mitre.org/cwe-6}'
      
      weaknesses = root[0]
      cwes = {}
  
      for w in weaknesses:
          _id = w.attrib['ID']
          _name = w.attrib['Name']
          _description = w.findall(namespace + 'Description')[0].text
          cves = []
          relationships = {}
          for relation in w.findall(namespace + 'Related_Weaknesses/' + namespace + 'Related_Weakness'):
              relationships[relation.attrib['CWE_ID']] = relation.attrib['Nature']
          for cwe in w.findall(namespace + 'Observed_Examples/' + namespace + 'Observed_Example/' + namespace + 'Reference'):
              cves.append(cwe.text)
          parents = [k for k, v in relationships.items() if v == 'ChildOf']
          cwes[_id] = {
              'id': _id,
              'name': _name,
              'description':  _description,
              'cve': cves,
              'parents': parents,
              'relationships': relationships
          }
      return cwes
  
  def read_cve(cve_file):
      """
      Parameters
      ----------
      cve_file : str
          The cve.xml file path.
  
      Returns
      -------
      cves : dict
          A CVE dictionary.
  
      """
      tree = ET.parse(cve_file)
      root = tree.getroot()
      
      namespace = '{http://cve.mitre.org/cve/downloads/1.0}'
      
      cves = {}
      
      for v in root:
          _id = v.attrib['name']
          cves[_id] = {
              'id': _id,
              'description': v.findall(namespace + 'desc')[0].text,
              'type': v.attrib['type']
          }
      return cves
      
  def add_children_graph(db):
      """
      Parameters
      ----------
      db : dict
  
      Returns
      -------
      db : dict
          Mutates the input dictionary with a map of parent-child hierarchy, 
          i.e. with a directed graph.
  
      """
      graph = {}
      for i in db:
          for r in db[i]['relationships']:
              if db[i]['relationships'][r] == 'ChildOf':
                  graph[r] = graph.get(r, [])
                  graph[r].append(i)
      for i in db:
          db[i]['children'] = graph[i] if i in graph else []
      return db
  
  def map_many_to_many(d, field):
      """
      Parameters
      ----------
      d : dict
          The input dictionary.
      field : str
          The filtering key.
  
      Returns
      -------
      m : dict
          A filtered dictionary composed only with non empty lists.
      m_inv : dict
          The inverted dictionary of m.
      """
      m = {k: v[field] for k, v in d.items() if len(v[field]) > 0}
      m_inv = {}
      for k, v in m.items():
          for i in v:
              m_inv[i] = m_inv.get(i, [])
              m_inv[i].append(k)
      return m, m_inv
      
      
  def start(input_dir, db_path):
      # read mapping
      capec = read_capec(input_dir + '/capec.xml')
      capec = add_children_graph(capec)
      
      cwe = read_cwe(input_dir + '/cwe.xml')
      cwe = add_children_graph(cwe)
      
      # cve = read_cve(input_dir + '/cve.xml')    
      
      # read suricata sid to cve mapping
      sid_cve = {}
      with open(input_dir +'/sid-cve', 'r') as f:
          for line in f:
             tmp = line.split(' ')
             sid_cve[tmp[0].strip(' "\'\t\r
  ')] = tmp[1].strip(' "\'\t\r
  ')
      
      cwe_cve, cve_cwe = map_many_to_many(cwe, 'cve')
      capec_cwe, cwe_capec = map_many_to_many(capec, 'cwe')
      
      capecdb = SqliteDict(db_path, autocommit=True)
      capecdb['capec'] = capec
      capecdb['cwe'] = cwe
      # capecdb['cve'] = cve
      capecdb['sid-cve'] = sid_cve
      capecdb['cve-cwe'] = cve_cwe
      capecdb['cwe-cve'] = cwe_cve
      capecdb['cwe-capec'] = cwe_capec
      capecdb['capec-cwe'] = capec_cwe
      capecdb.close()
  
  
  if len(argv) - 1 == 2 and path.exists(argv[1]):
      start(argv[1], argv[2])
  else:
      print('Error, please provide the paths of eve.json and capecdb')