import os
import tempfile
import pytest
from ayn_antivirus.signatures.db.hash_db import HashDatabase
from ayn_antivirus.signatures.db.ioc_db import IOCDatabase

def test_hash_db_create(tmp_path):
    db = HashDatabase(str(tmp_path / "test.db"))
    db.initialize()
    assert db.count() == 0
    db.close()

def test_hash_db_add_and_lookup(tmp_path):
    db = HashDatabase(str(tmp_path / "test.db"))
    db.initialize()
    db.add_hash("abc123hash", "TestMalware", "virus", "high", "test")
    result = db.lookup("abc123hash")
    assert result is not None
    assert result["threat_name"] == "TestMalware"
    db.close()

def test_hash_db_bulk_add(tmp_path):
    db = HashDatabase(str(tmp_path / "test.db"))
    db.initialize()
    records = [
        ("hash1", "Malware1", "virus", "high", "test", ""),
        ("hash2", "Malware2", "malware", "medium", "test", ""),
        ("hash3", "Miner1", "miner", "high", "test", ""),
    ]
    count = db.bulk_add(records)
    assert count == 3
    assert db.count() == 3
    db.close()

def test_ioc_db_ips(tmp_path):
    db = IOCDatabase(str(tmp_path / "test.db"))
    db.initialize()
    db.add_ip("1.2.3.4", "BotnetC2", "c2", "feodo")
    result = db.lookup_ip("1.2.3.4")
    assert result is not None
    ips = db.get_all_malicious_ips()
    assert "1.2.3.4" in ips
    db.close()

def test_ioc_db_domains(tmp_path):
    db = IOCDatabase(str(tmp_path / "test.db"))
    db.initialize()
    db.add_domain("evil.com", "Phishing", "phishing", "threatfox")
    result = db.lookup_domain("evil.com")
    assert result is not None
    domains = db.get_all_malicious_domains()
    assert "evil.com" in domains
    db.close()