import json
import os
import re
import sys
import tldextract

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, IOC
from malware_data_science.utils import read_json

db = DB()

def new_ioc(attribute_type_id: int, data: str):
  new_ioc = IOC()
  new_ioc.attribute_type_id = attribute_type_id
  new_ioc.data = data
  return new_ioc

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

def seed_iocs(attribute_type_id: int, ioc_key: str, ioc_sub_key: str):
  for malware in Malware.all():
    iocs = iocs_from_json(malware.sha256)
    if iocs == None:
      continue
    for entry in iocs[ioc_key]:
      if IOC.find_one(entry[ioc_sub_key], malware.id, attribute_type_id):
        continue
      malware.iocs.append(new_ioc(attribute_type_id, entry[ioc_sub_key]))
    db.session.commit()

def seed_file_hashes(attribute_type_id: int, ioc_key: str, ioc_sub_key: str, ioc_sub_key2: str):
  for malware in Malware.all():
    iocs = iocs_from_json(malware.sha256)
    if iocs == None:
      continue
    for entry in iocs[ioc_key]:
      for subentry in entry[ioc_sub_key]:
        data = subentry[ioc_sub_key2]
        if IOC.find_one(data, malware.id, attribute_type_id):
          continue
        malware.iocs.append(new_ioc(attribute_type_id, data))
    db.session.commit()

def seed_process_attributes(attribute_type_id: int, ioc_key: str, ioc_sub_key: str):
  for malware in Malware.all():
    iocs = iocs_from_json(malware.sha256)
    if iocs == None:
      continue
    for entry in iocs[ioc_key]:
      for data in entry[ioc_sub_key]:
        if IOC.find_one(data, malware.id, attribute_type_id):
          continue
        malware.iocs.append(new_ioc(attribute_type_id, data))
    db.session.commit()

def seed_domains():
  attribute_type_id = 1
  for malware in Malware.all():
    iocs = iocs_from_json(malware.sha256)
    if iocs == None:
      continue
    for entry in iocs['domains']:
      ext = tldextract.extract(entry['domain'])
      if ext.registered_domain:
        if IOC.find_one(ext.registered_domain, malware.id, attribute_type_id):
          continue
        malware.iocs.append(new_ioc(attribute_type_id, ext.registered_domain))
    db.session.commit()

seed_iocs(2, 'email_addresses', 'email_address')
seed_iocs(3, 'emails', 'email')
seed_iocs(4, 'filenames', 'filename')
seed_iocs(9, 'ips', 'ip_address')
seed_iocs(10, 'mutexes', 'mutex_name')
seed_iocs(13, 'urls', 'url')

seed_file_hashes(5, 'files', 'hashes', 'sha1_hash')
seed_file_hashes(6, 'files', 'hashes', 'md5_hash')
seed_file_hashes(7, 'files', 'hashes', 'sha256_hash')
seed_file_hashes(8, 'files', 'hashes', 'ssdeep_hash')

seed_process_attributes(11, 'processes', 'process_names')
seed_process_attributes(14, 'processes', 'threat_names')

seed_domains()

# TODO: Add IOCs related to domains and registry keys
# seed_iocs(12, 'registry', 'registry_key')
