"""
Database models for the Malware Website Crawler
"""
from sqlalchemy import Column, Integer, String, Text, DateTime, Float, Boolean, ForeignKey, JSON
from sqlalchemy.orm import relationship
from sqlalchemy.sql import func
from database import Base

class ScanJob(Base):
    __tablename__ = 'scan_jobs'
    
    id = Column(Integer, primary_key=True)
    target_url = Column(String(500), nullable=False)
    scan_type = Column(String(50), default='website')
    status = Column(String(20), default='pending')  # pending, running, completed, failed
    task_id = Column(String(100))
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    completed_at = Column(DateTime(timezone=True))
    
    # Relationships
    findings = relationship("MalwareFinding", back_populates="scan_job")

class MalwareFinding(Base):
    __tablename__ = 'malware_findings'
    
    id = Column(Integer, primary_key=True)
    scan_job_id = Column(Integer, ForeignKey('scan_jobs.id'), nullable=False)
    filename = Column(String(255), nullable=False)
    file_path = Column(String(500))
    file_hash = Column(String(64), nullable=False)  # SHA256
    file_size = Column(Integer)
    file_type = Column(String(100))
    
    # Analysis results
    malware_name = Column(String(255))
    confidence = Column(Float, default=0.0)
    features = Column(JSON)  # List of detected features
    ip_addresses = Column(JSON)  # List of discovered IP addresses
    attack_vectors = Column(JSON)  # List of attack vectors
    context_words = Column(JSON)  # List of contextual words
    
    # External analysis
    virustotal_detected = Column(Boolean, default=False)
    virustotal_engines = Column(Integer, default=0)
    known_malware_id = Column(Integer, ForeignKey('known_malware.id'))
    
    # Metadata
    discovered_at = Column(DateTime(timezone=True), server_default=func.now())
    
    # Relationships
    scan_job = relationship("ScanJob", back_populates="findings")
    known_malware = relationship("KnownMalware", back_populates="findings")

class KnownMalware(Base):
    __tablename__ = 'known_malware'
    
    id = Column(Integer, primary_key=True)
    name = Column(String(255), nullable=False)
    family = Column(String(100))
    description = Column(Text)
    
    # File signatures
    md5_hash = Column(String(32))
    sha1_hash = Column(String(40))
    sha256_hash = Column(String(64))
    
    # Characteristics
    features = Column(JSON)  # Common features
    attack_vectors = Column(JSON)  # Known attack vectors
    c2_servers = Column(JSON)  # Known C2 servers
    
    # Metadata
    first_seen = Column(DateTime(timezone=True))
    last_updated = Column(DateTime(timezone=True), server_default=func.now())
    source = Column(String(100))  # Where this signature came from
    
    # Relationships
    findings = relationship("MalwareFinding", back_populates="known_malware")

class ContextWord(Base):
    __tablename__ = 'context_words'
    
    id = Column(Integer, primary_key=True)
    word = Column(String(100), unique=True, nullable=False)
    category = Column(String(50))  # e.g., 'business', 'technology', 'finance'
    frequency = Column(Integer, default=1)
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    last_used = Column(DateTime(timezone=True), server_default=func.now())
