import networkx as nx
import itertools
import json
import re

from pyvis.network import Network
from malware_data_science.db import DB, Malware, Campaign, CampaignIOC
from networkx.drawing.nx_agraph import write_dot

class NetworkAttributeAnalyzer:
  def __init__(self, attribute_type_id: int, ioc_color: str):
    self.db = DB()
    self.G = nx.Graph()
    self.attribute_type_id = attribute_type_id
    self.ioc_color = ioc_color
    self.shortest_path = False
    self.output = 'html'

  def build_graph(self) -> None:
    for malware in Malware.all():
      if malware.count_by_attribute_type_id(self.attribute_type_id) == 0:
        continue
      if self.output == 'html':
        self.G.add_node(malware.sha256, label=malware.threat_names(), type='malware', group=1, size=30, color='#8A0808')
      else:
        self.G.add_node(malware.sha256, label='M', type='malware', fontsize=10.0, fontcolor='white', fontname='Helvetica Bold',
                        shape='circle', color='white', style='filled,bold', fillcolor='red4',
                        height=0.5, width=0.5)
      for ioc in malware.ioc_by_attribute_type_id(self.attribute_type_id):
        if self.output == 'html':
          self.G.add_node(ioc.data, label=ioc.data, type='ioc', group=2, size=25, color=self.ioc_color)
          self.G.add_edge(malware.sha256, ioc.data, color='#D8D8D8')
        else:
          self.G.add_node(ioc.data, type='ioc', label='',
                          shape='circle', color='white', style='filled,bold', fillcolor=self.ioc_color, fontcolor=self.ioc_color,
                          height=0.4, width=0.4)
          self.G.add_edge(malware.sha256, ioc.data, color='gray77', style='filled,bold')

  def malware_nodes(self, g) -> set:
    return set(node for node, attr in g.nodes(data=True) if attr['type'] == 'malware')
  
  def remove_isolated_nodes(self) -> None:
    self.G.remove_nodes_from(list(nx.isolates(self.G)))
    not_required_nodes = []
    for c in nx.connected_components(self.G):
      g = self.G.subgraph(c)
      if len(self.malware_nodes(g)) > 1:
        if self.shortest_path == True:
          shortest_paths =[]
          for source, target in itertools.combinations(self.malware_nodes(g),2):
            for path in nx.shortest_path(self.G, source=source, target=target):
              shortest_paths.append(path)
            for ioc in set(g) - self.malware_nodes(g):
              if len(list(g.neighbors(ioc))) < 2 or ioc not in shortest_paths:
                not_required_nodes.append(ioc)
        else:
          for ioc in set(g) - self.malware_nodes(g):
            if len(list(g.neighbors(ioc))) < 2:
              not_required_nodes.append(ioc)
      else:
        for node in g.nodes:
          not_required_nodes.append(node)
    self.G.remove_nodes_from(not_required_nodes)

  def enumerate_subgraphs(self) -> None:
    campaign_number = 1
    for c in nx.connected_components(self.G):
      g = self.G.subgraph(c)
      nx.set_node_attributes(g, str(campaign_number), 'campaign_number')
      if self.output != 'html':
        nx.set_node_attributes(g, str(campaign_number), 'label')
      campaign_number += 1

  def save_campaign(self) -> None:
    self.enumerate_subgraphs()
    for c in nx.connected_components(self.G):
      g = self.G.subgraph(c)
      for node, attr in g.nodes(data=True):
        if attr['type'] == 'malware':
          malware = Malware.one_by_sha256(node)
          campaign = Campaign()
          campaign.malware = malware
          campaign.number = attr['campaign_number']
          campaign.attribute_type_id = self.attribute_type_id
          for ioc_data in g.neighbors(node):
            campaign_ioc = CampaignIOC()
            campaign_ioc.ioc_data = ioc_data
            campaign.campaign_iocs.append(campaign_ioc)
          malware.campaigns.append(campaign)
          malware.session.commit()
  
  def delete_campaign(self) -> None:
    Campaign.delete_by_attribute_type_id(self.attribute_type_id)
    self.db.session.commit()

  def run_analysis(self) -> None:
    self.build_graph()
    self.remove_isolated_nodes()
    self.delete_campaign()
    self.save_campaign()

  def save_graph(self, dest_path: str) -> None:
    if self.output == 'html':
      nt = Network(height='100%', width='100%', font_color='#848484')
      nt.from_nx(self.G)
      nt.save_graph(dest_path)
    else:
      write_dot(self.G, dest_path)
