"""
Script to load known malware signatures into the database
"""
import json
import sys
import os
from datetime import datetime

# Add parent directory to path to import modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from database import init_db, db_session
from models import KnownMalware

def load_known_malware():
    """Load known malware signatures from JSON file"""
    try:
        # Initialize database
        init_db()
        
        # Load malware signatures from JSON file
        with open('known_malware.json', 'r') as f:
            data = json.load(f)
        
        malware_signatures = data.get('malware_signatures', [])
        
        print(f"Loading {len(malware_signatures)} malware signatures...")
        
        for signature in malware_signatures:
            # Check if malware already exists
            existing = db_session.query(KnownMalware).filter(
                KnownMalware.name == signature['name']
            ).first()
            
            if existing:
                print(f"Updating existing signature: {signature['name']}")
                # Update existing record
                existing.family = signature.get('family')
                existing.description = signature.get('description')
                existing.md5_hash = signature.get('md5_hash')
                existing.sha1_hash = signature.get('sha1_hash')
                existing.sha256_hash = signature.get('sha256_hash')
                existing.features = signature.get('features', [])
                existing.attack_vectors = signature.get('attack_vectors', [])
                existing.c2_servers = signature.get('c2_servers', [])
                existing.source = signature.get('source')
                existing.last_updated = datetime.utcnow()
            else:
                print(f"Adding new signature: {signature['name']}")
                # Create new record
                malware = KnownMalware(
                    name=signature['name'],
                    family=signature.get('family'),
                    description=signature.get('description'),
                    md5_hash=signature.get('md5_hash'),
                    sha1_hash=signature.get('sha1_hash'),
                    sha256_hash=signature.get('sha256_hash'),
                    features=signature.get('features', []),
                    attack_vectors=signature.get('attack_vectors', []),
                    c2_servers=signature.get('c2_servers', []),
                    first_seen=datetime.fromisoformat(signature.get('first_seen', datetime.utcnow().isoformat())),
                    source=signature.get('source')
                )
                db_session.add(malware)
        
        # Commit changes
        db_session.commit()
        print("Successfully loaded malware signatures!")
        
    except Exception as e:
        print(f"Error loading malware signatures: {str(e)}")
        db_session.rollback()
        sys.exit(1)
    finally:
        db_session.close()

if __name__ == '__main__':
    load_known_malware()
