import json
import os
import re
import sys

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

from tqdm import tqdm
from vmray.rest_api import VMRayRESTAPI, VMRayRESTAPIError
from malware_data_science.config import Config
from malware_data_science.db import Malware, DB

class MetadataDowloader(Config):
  def __init__(self):
    self.db = DB()
    super().__init__()

  def download(self, desc: str, endpoint: str):
    api = VMRayRESTAPI(self.vmray['server'], self.vmray['api_key'])
    counter = 0
    with tqdm(total=Malware.count(), desc='Downloading ' + desc) as pbar:
      for malware in Malware.all():
        try:
          data = api.call('GET', '/rest/sample/' + str(malware.sample_id) + endpoint)
          self.save_as_json(malware.sha256, data, re.sub('^\/','', endpoint))
        except VMRayRESTAPIError:
          continue
        counter += 1
        pbar.update(counter)

  def save_as_json(self, filename: str, content: dict, subdir: str):
    dest_dir = os.path.join(self.data_dir, subdir)
    os.makedirs(dest_dir, exist_ok = True)
    file_path = os.path.join(dest_dir, filename + '.json')
    with open(file_path, 'w') as file:
      json.dump(content, file, indent=4)

downloader = MetadataDowloader()
downloader.download('IOCS', '/iocs')
downloader.download('Mitre Attack data', '/mitre_attack')
