#!/usr/bin/python
import argparse
import networkx as nx
import pydot
import pathlib

from networkx.drawing.nx_agraph import write_dot
from networkx.algorithms import bipartite
from pathlib import Path
from malware_data_science.pe_util import PEUtil
from malware_data_science.pe_resource import PEResource

class NetworkResourceAnalyzer(PEUtil):
  def __init__(self, target_path):
    self.target_path = target_path
    self.G = nx.Graph()
    self.malware_type = 'malware'
    self.image_type = 'image'

  def build_graph(self):
    i = 0
    for fullpath in self.pe_files(self.target_path):
      sample_name = fullpath.split('/')[-1]
      malware_resource = PEResource(fullpath)
      malware_resource.extract_images()
      malware_resource.image_hashes()
      i += 1
      malware_name = (f'M{i}')
      self.G.add_node(sample_name, label=malware_name, type=self.malware_type,
                      fontname='helvetica bold', fontsize=15.0, fontcolor='red4',
                      shape='circle', color='white', style='filled,bold', fillcolor='red4',
                      height=1, width=1)
      for _img_path, image_hash in malware_resource.images:
        if not image_hash in self.G:
          self.G.add_node(image_hash, label='', type=self.image_type,
                          shape='circle', color='white', style='filled,bold', fillcolor='darkslategrey',
                          height=0.6, width=0.6)
        self.G.add_edge(sample_name, image_hash, color='gray77', style='filled,bold')

  def remove_isolated_nodes(self):
    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:
        for r in set(g) - self.malware_nodes(g):
          if len(list(g.neighbors(r))) < 2:
            not_required_nodes.append(r)
      else:
        for node in g.nodes:
          not_required_nodes.append(node)
    self.G.remove_nodes_from(not_required_nodes)
  
  def save_network(self, file_path):
    write_dot(self.G, file_path)
  
  def malware_nodes(self, g):
    return set(node for node, attr in g.nodes(data=True) if attr['type'] == 'malware')

  def save_malware_nodes(self, file_path):
    malware_network = bipartite.projected_graph(self.G, self.malware_nodes(self.G))
    write_dot(malware_network, file_path)
  
  def save_image_nodes(self, file_path):
    resource = set(self.G) - self.malware_nodes(self.G)
    resource_network = bipartite.projected_graph(self.G, resource)
    write_dot(resource_network, file_path)