54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
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()
|