import os
import re

from datetime import datetime
from sqlalchemy import Column, ForeignKey, Integer, String, DateTime, UniqueConstraint, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from sqlalchemy.orm.exc import NoResultFound
from malware_data_science.config import Config

Base = declarative_base()

class Malware(Base):
  __tablename__ = 'malware'
  id = Column(Integer, primary_key=True, unique=True)
  name = Column(String())
  md5 = Column(String(32), nullable=False, unique=True)
  sha256 = Column(String(64), nullable=False, unique=True)
  ssdeep_hash = Column(String(64))
  created_at = Column(DateTime(), nullable=False)
  sample_id = Column(Integer, nullable=True, unique=True)
  sample_type = Column(String())
  iocs = relationship('IOC', lazy='dynamic')
  campaigns = relationship('Campaign')
  malware_threat_indicators = relationship('MalwareThreatIndicator')
  threat_indicators = relationship('ThreatIndicator', secondary='malware_threat_indicators')

  @classmethod
  def one_by_sha256(cls, sha256: str):
    try:
      return Base.session.query(Malware).filter(Malware.sha256 == sha256).one()
    except NoResultFound:
      return None
  
  @classmethod
  def all(cls):
    try:
      return Base.session.query(Malware).all()
    except NoResultFound:
      return None
  
  @classmethod
  def count(cls) -> int:
    try:
      return Base.session.query(Malware).count()
    except NoResultFound:
      return 0

  def count_by_attribute_type_id(self, attribute_type_id: int) -> int:
    return self.iocs.filter(IOC.attribute_type_id == attribute_type_id).count()
 
  def ioc_by_attribute_type_id(self, attribute_type_id: int):
    return self.iocs.filter(IOC.attribute_type_id == attribute_type_id).all()

  def threat_names(self):
    threat_names = []
    for ioc in self.iocs.filter(IOC.attribute_type_id == 14).limit(3).all():
      threat_name = re.sub(r'\.\w*$', '', ioc.data.replace('Gen:Variant.', ''))
      if not threat_name in threat_names:
        threat_names.append(threat_name)

    if len(threat_names) > 0:
      return '/'.join(threat_names)
    else:
      return ''

class AttributeType(Base):
  __tablename__ = 'attribute_types'
  id = Column(Integer, primary_key=True, unique=True)
  name = Column(String(), nullable=False, unique=True)
  iocs = relationship('IOC')

  @classmethod
  def one_by_name(cls, name: str):
    try:
      return Base.session.query(AttributeType).filter(AttributeType.name == name).one()
    except NoResultFound:
      return None

class IOC(Base):
  __tablename__ = 'iocs'
  id = Column(Integer, primary_key=True, unique=True)
  __table_args__ = (
    UniqueConstraint(
      'malware_id',
      'data',
      'attribute_type_id',
      name='malware_id_data_attribute_type_id_unique_key'
    ),
  )
  id = Column(Integer, primary_key=True, unique=True)
  malware_id = Column(Integer, ForeignKey('malware.id'))
  attribute_type_id = Column(Integer, ForeignKey('attribute_types.id'))
  data = Column(String())
  malware = relationship('Malware')
  attribute_type = relationship('AttributeType')

  @classmethod
  def find_one(cls, data: str, malware_id: int, attribute_type_id: int):
    try:
      return Base.session.query(IOC).\
        filter(IOC.data == data). \
        filter(IOC.malware_id == malware_id). \
        filter(IOC.attribute_type_id == attribute_type_id). \
        one()
    except NoResultFound:
      return None

class Campaign(Base):
  __tablename__ = 'campaigns'
  __table_args__ = (
    UniqueConstraint('malware_id', 'attribute_type_id', 'number', name='malware_id_attribute_type_id_number_unique_key'),
  )
  id = Column(Integer, primary_key=True, unique=True)
  malware_id = Column(Integer, ForeignKey('malware.id'))
  attribute_type_id = Column(Integer, ForeignKey('attribute_types.id'))
  number = Column(Integer, nullable=False)
  malware = relationship('Malware')
  attribute_type = relationship('AttributeType')
  campaign_iocs = relationship('CampaignIOC', cascade='all, delete, delete-orphan', passive_deletes=True)

  @classmethod
  def all_by_attribute_type_id(cls, attribute_type_id: int):
    return Base.session.query(Campaign). \
            filter(Campaign.attribute_type_id == attribute_type_id). \
            all()

  @classmethod
  def all_by_attribute_id_created_at_desc(cls, attribute_type_id: int):
    try:
      return Base.session.query(Campaign). \
              join(Malware). \
              filter(Campaign.attribute_type_id == attribute_type_id). \
              order_by(Malware.created_at.desc()). \
              all()
    except NoResultFound:
      return None

  @classmethod
  def delete_by_attribute_type_id(cls, attribute_type_id: int) -> None:
    # TODO: Fix DELETE ON CASCADE
    for campaign in Base.session.query(Campaign).filter(Campaign.attribute_type_id == attribute_type_id).all():
      Base.session.query(CampaignIOC).filter(CampaignIOC.campaign_id == campaign.id).delete()
    Base.session.query(Campaign).filter(Campaign.attribute_type_id == attribute_type_id).delete()  

class CampaignIOC(Base):
  __tablename__ = 'campaign_iocs'
  __table_args__ = (
    UniqueConstraint('campaign_id', 'ioc_data', name='malware_id_ioc_data_unique_key'),
  )
  id = Column(Integer, primary_key=True, unique=True)
  campaign_id = Column(Integer, ForeignKey('campaigns.id', ondelete='CASCADE'), nullable=False)
  ioc_data = Column(Integer, nullable=False)
  campaign = relationship('Campaign')

class AttackTactic(Base):
  __tablename__ = 'attack_tactics'
  id = Column(Integer, primary_key=True, unique=True)
  name = Column(String(), nullable=False, unique=True)
  attack_techniques = relationship('AttackTechnique')

  @classmethod
  def all(cls):
    try:
      return Base.session.query(AttackTactic).all()
    except NoResultFound:
      return None

  @classmethod
  def one_by_name(cls, name: str):
    try:
      return Base.session.query(AttackTactic).filter(AttackTactic.name == name).one()
    except NoResultFound:
      return None

class AttackTechnique(Base):
  __tablename__ = 'attack_techniques'
  __table_args__ = (
    UniqueConstraint('name', 'technique_id', 'attack_tactic_id', name='name_technique_id_attack_tactic_id_unique_key'),
  )
  id = Column(Integer, primary_key=True, unique=True)
  name = Column(String(), nullable=False)
  technique_id = Column(String(), nullable=False)
  attack_tactic_id = Column(Integer, ForeignKey('attack_tactics.id'))
  attack_tactic = relationship('AttackTactic')
  threat_indicators = relationship('ThreatIndicator')

  def name_and_technique_id(self) -> str:
    return ' - '.join([self.name, self.technique_id])

  @classmethod
  def one_by_name_technique_id(cls, name: str, technique_id: str, attack_tactic_id: int):
    try:
      return Base.session.query(AttackTechnique). \
        filter(AttackTechnique.name == name). \
        filter(AttackTechnique.technique_id == technique_id). \
        filter(AttackTechnique.attack_tactic_id == attack_tactic_id). \
        one()
    except NoResultFound:
      return None

class ThreatIndicatorCategory(Base):
  __tablename__ = 'threat_indicator_categories'
  id = Column(Integer, primary_key=True, unique=True)
  name = Column(String(), nullable=False, unique=True)
  threat_indicators = relationship('ThreatIndicator')

  @classmethod
  def one_by_name(cls, name: str):
    try:
      return Base.session.query(ThreatIndicatorCategory).filter(ThreatIndicatorCategory.name == name).one()
    except NoResultFound:
      return None

class ThreatIndicator(Base):
  __tablename__ = 'threat_indicators'
  __table_args__ = (
    UniqueConstraint('operation', 'attack_technique_id', 'threat_indicator_category_id', name='operation_attack_technique_id_threat_indicator_category_id_unique_key'),
  )
  id = Column(Integer, primary_key=True, unique=True)
  operation = Column(String(), nullable=False)
  threat_indicator_category_id = Column(Integer, ForeignKey('threat_indicator_categories.id'))
  attack_technique_id = Column(Integer, ForeignKey('attack_techniques.id'))
  threat_indicator_category = relationship('ThreatIndicatorCategory')
  attack_technique = relationship('AttackTechnique')
  #malwares = relationship('MalwareThreatIndicator', secondary='malware_threat_indicators')

  @classmethod
  def one_by_operation_attack_technique_id(cls, operation: str, attack_technique_id: int, threat_indicator_category_id: int):
      try:
        return Base.session.query(ThreatIndicator). \
          filter(ThreatIndicator.operation == operation). \
          filter(ThreatIndicator.attack_technique_id == attack_technique_id). \
          filter(ThreatIndicator.threat_indicator_category_id == threat_indicator_category_id). \
          one()
      except NoResultFound:
        return None
  
  def operation_and_category_name(self) -> str:
    return ' - '.join([self.operation, self.threat_indicator_category.name])

class MalwareThreatIndicator(Base):
  __tablename__ = 'malware_threat_indicators'
  id = Column(Integer, primary_key=True, unique=True)
  __table_args__ = (
    UniqueConstraint('score', 'malware_id', 'threat_indicator_id', name='score_malware_id_threat_indicator_id_unique_key'),
  )
  score = Column(Integer(), nullable=False)
  malware_id = Column(Integer, ForeignKey('malware.id'))
  threat_indicator_id = Column(Integer, ForeignKey('threat_indicators.id'))
  malware = relationship('Malware')
  threat_indicator = relationship('ThreatIndicator')

  @classmethod
  def one_by_malware_id_threat_indicator_id(cls, malware_id: int, threat_indicator_id: int, score: int):
    try:
      return Base.session.query(MalwareThreatIndicator). \
        filter(MalwareThreatIndicator.malware_id == malware_id). \
        filter(MalwareThreatIndicator.threat_indicator_id == threat_indicator_id). \
        filter(MalwareThreatIndicator.score == score). \
        one()
    except NoResultFound:
      return None

class Classification(Base):
  __tablename__ = 'classification'
  id = Column(Integer, primary_key=True, unique=True)
  name = Column(String(), nullable=False, unique=True)
  malwares = relationship('Malware', secondary='malware_classification')

  @classmethod
  def one_by_name(cls, name: str):
    try:
      return Base.session.query(Classification).filter(Classification.name == name).one()
    except NoResultFound:
      return None

class MalwareClassification(Base):
  __tablename__ = 'malware_classification'
  id = Column(Integer, primary_key=True, unique=True)
  __table_args__ = (
    UniqueConstraint('malware_id', 'classification_id', name='malware_id_classification_id_unique_key'),
  )
  malware_id = Column(Integer, ForeignKey('malware.id'))
  classification_id = Column(Integer, ForeignKey('classification.id'))
  malware = relationship('Malware')
  classification = relationship('Classification')

class DB(Config):
  def __init__(self):
    self.config = Config()
    db_path = 'sqlite:///' + self.config.database['path']
    self.engine = create_engine(db_path)
    self.session = sessionmaker(bind=self.engine)()
    Base.metadata.create_all(bind=self.engine)
    Base.metadata.extend_existing = True
    Base.session = self.session
