import json
import os
import re
import sys

sys.path.append(os.path.dirname(os.path.realpath(__file__))+'/../..')

from pathlib import Path
from malware_data_science.config import Config
from malware_data_science.db import DB, Malware, AttackTactic, AttackTechnique, \
                                    ThreatIndicatorCategory, ThreatIndicator, \
                                    MalwareThreatIndicator
from malware_data_science.utils import read_json

def attack_from_json(filename: str):
  config = Config()
  src_dir = os.path.join(config.data_dir, 'mitre_attack')
  json_path = os.path.join(src_dir, filename + '.json')
  json_data = read_json(json_path)
  if json_data != None:
    return json_data['mitre_attack_techniques']
  else:
    return None

db = DB()
for malware in Malware.all():
  entries = attack_from_json(malware.sha256)
  if entries == None or len(entries) == 0:
    continue
  
  for entry in entries:
    for tactic_name in entry['tactics']:
      attack_tactic = AttackTactic.one_by_name(tactic_name)
      if attack_tactic == None:
        attack_tactic = AttackTactic()
        attack_tactic.name = tactic_name
        db.session.add(attack_tactic)
        db.session.commit()

      attack_technique = AttackTechnique.one_by_name_technique_id(entry['technique'], entry['technique_id'], attack_tactic.id)
      if attack_technique == None:
        attack_technique = AttackTechnique()
        attack_technique.name = entry['technique']
        attack_technique.technique_id = entry['technique_id']
        attack_technique.attack_tactic_id = attack_tactic.id 
        db.session.add(attack_technique)
      else:
        attack_tactic.attack_techniques.append(attack_technique)
      db.session.commit()

      for entry_ti in entry['threat_indicators']:
        threat_indicator_category = ThreatIndicatorCategory.one_by_name(entry_ti['category'])
        if threat_indicator_category == None:
          threat_indicator_category = ThreatIndicatorCategory()
          threat_indicator_category.name = entry_ti['category']
          db.session.add(threat_indicator_category)
          db.session.commit()
        
        threat_indicator = ThreatIndicator.one_by_operation_attack_technique_id(entry_ti['operation'], attack_technique.id, threat_indicator_category.id)
        if threat_indicator == None:
          threat_indicator = ThreatIndicator()
          threat_indicator.operation = entry_ti['operation']
          threat_indicator.attack_technique_id = attack_technique.id
          threat_indicator.threat_indicator_category_id = threat_indicator_category.id
          db.session.add(threat_indicator)
          db.session.commit()
        
        malware_threat_indicator = MalwareThreatIndicator.one_by_malware_id_threat_indicator_id(malware.id, threat_indicator.id, int(entry_ti['score']))
        if malware_threat_indicator == None:
          malware_threat_indicator = MalwareThreatIndicator()
          malware_threat_indicator.malware_id = malware.id
          malware_threat_indicator.threat_indicator_id = threat_indicator.id
          malware_threat_indicator.score = int(entry_ti['score'])
          db.session.add(malware_threat_indicator)
          db.session.commit()
