remove infra.md.example, infra.md is the source of truth
This commit is contained in:
Azreen Jamal
192
ayn-antivirus/ayn_antivirus/detectors/signature_detector.py
Normal file
192
ayn-antivirus/ayn_antivirus/detectors/signature_detector.py
Normal file
@@ -0,0 +1,192 @@
|
||||
"""AYN Antivirus — Signature-based Detector.
|
||||
|
||||
Looks up file hashes against the threat signature database populated by
|
||||
the feed update pipeline (MalwareBazaar, ThreatFox, etc.). Uses
|
||||
:class:`~ayn_antivirus.signatures.db.hash_db.HashDatabase` so that
|
||||
definitions written by ``ayn-antivirus update`` are immediately available
|
||||
for detection.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from ayn_antivirus.constants import DEFAULT_DB_PATH
|
||||
from ayn_antivirus.detectors.base import BaseDetector, DetectionResult
|
||||
from ayn_antivirus.utils.helpers import hash_file as _hash_file_util
|
||||
|
||||
logger = logging.getLogger("ayn_antivirus.detectors.signature")
|
||||
|
||||
_VALID_SEVERITIES = {"CRITICAL", "HIGH", "MEDIUM", "LOW"}
|
||||
|
||||
|
||||
class SignatureDetector(BaseDetector):
|
||||
"""Detect known malware by matching file hashes against the signature DB.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
db_path:
|
||||
Path to the shared SQLite database that holds the ``threats``,
|
||||
``ioc_ips``, ``ioc_domains``, and ``ioc_urls`` tables.
|
||||
"""
|
||||
|
||||
def __init__(self, db_path: str | Path = DEFAULT_DB_PATH) -> None:
|
||||
self.db_path = str(db_path)
|
||||
self._hash_db = None
|
||||
self._ioc_db = None
|
||||
self._loaded = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# BaseDetector interface
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return "signature_detector"
|
||||
|
||||
@property
|
||||
def description(self) -> str:
|
||||
return "Hash-based signature detection using threat intelligence feeds"
|
||||
|
||||
def detect(
|
||||
self,
|
||||
file_path: str | Path,
|
||||
file_content: Optional[bytes] = None,
|
||||
file_hash: Optional[str] = None,
|
||||
) -> List[DetectionResult]:
|
||||
"""Check the file's hash against the ``threats`` table.
|
||||
|
||||
If *file_hash* is not supplied it is computed on the fly.
|
||||
"""
|
||||
self._ensure_loaded()
|
||||
results: List[DetectionResult] = []
|
||||
|
||||
if not self._hash_db:
|
||||
return results
|
||||
|
||||
# Compute hash if not provided.
|
||||
if not file_hash:
|
||||
try:
|
||||
file_hash = _hash_file_util(str(file_path), algo="sha256")
|
||||
except Exception:
|
||||
return results
|
||||
|
||||
# Also compute MD5 for VirusShare lookups.
|
||||
md5_hash = None
|
||||
try:
|
||||
md5_hash = _hash_file_util(str(file_path), algo="md5")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Look up SHA256 first, then MD5.
|
||||
threat = self._hash_db.lookup(file_hash)
|
||||
if not threat and md5_hash:
|
||||
threat = self._hash_db.lookup(md5_hash)
|
||||
if threat:
|
||||
severity = (threat.get("severity") or "HIGH").upper()
|
||||
if severity not in _VALID_SEVERITIES:
|
||||
severity = "HIGH"
|
||||
results.append(DetectionResult(
|
||||
threat_name=threat.get("threat_name", "Malware.Known"),
|
||||
threat_type=threat.get("threat_type", "MALWARE"),
|
||||
severity=severity,
|
||||
confidence=100,
|
||||
details=(
|
||||
f"Known threat signature match "
|
||||
f"(source: {threat.get('source', 'unknown')}). "
|
||||
f"Hash: {file_hash[:16]}... "
|
||||
f"Details: {threat.get('details', '')}"
|
||||
),
|
||||
detector_name=self.name,
|
||||
))
|
||||
|
||||
return results
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# IOC lookup helpers (used by engine for network enrichment)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def lookup_hash(self, file_hash: str) -> Optional[Dict]:
|
||||
"""Look up a single hash. Returns threat info dict or ``None``."""
|
||||
self._ensure_loaded()
|
||||
if not self._hash_db:
|
||||
return None
|
||||
return self._hash_db.lookup(file_hash)
|
||||
|
||||
def lookup_ip(self, ip: str) -> Optional[Dict]:
|
||||
"""Look up an IP against the IOC database."""
|
||||
self._ensure_loaded()
|
||||
if not self._ioc_db:
|
||||
return None
|
||||
return self._ioc_db.lookup_ip(ip)
|
||||
|
||||
def lookup_domain(self, domain: str) -> Optional[Dict]:
|
||||
"""Look up a domain against the IOC database."""
|
||||
self._ensure_loaded()
|
||||
if not self._ioc_db:
|
||||
return None
|
||||
return self._ioc_db.lookup_domain(domain)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Statistics
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_stats(self) -> Dict:
|
||||
"""Return signature / IOC database statistics."""
|
||||
self._ensure_loaded()
|
||||
stats: Dict = {"hash_count": 0, "loaded": self._loaded}
|
||||
if self._hash_db:
|
||||
stats["hash_count"] = self._hash_db.count()
|
||||
stats.update(self._hash_db.get_stats())
|
||||
if self._ioc_db:
|
||||
stats["ioc_ips"] = len(self._ioc_db.get_all_malicious_ips())
|
||||
stats["ioc_domains"] = len(self._ioc_db.get_all_malicious_domains())
|
||||
return stats
|
||||
|
||||
@property
|
||||
def signature_count(self) -> int:
|
||||
"""Number of hash signatures currently loaded."""
|
||||
self._ensure_loaded()
|
||||
return self._hash_db.count() if self._hash_db else 0
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close database connections."""
|
||||
if self._hash_db:
|
||||
self._hash_db.close()
|
||||
self._hash_db = None
|
||||
if self._ioc_db:
|
||||
self._ioc_db.close()
|
||||
self._ioc_db = None
|
||||
self._loaded = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _ensure_loaded(self) -> None:
|
||||
"""Lazy-load the database connections on first use."""
|
||||
if self._loaded:
|
||||
return
|
||||
if not self.db_path:
|
||||
logger.warning("No signature DB path configured")
|
||||
self._loaded = True
|
||||
return
|
||||
try:
|
||||
from ayn_antivirus.signatures.db.hash_db import HashDatabase
|
||||
from ayn_antivirus.signatures.db.ioc_db import IOCDatabase
|
||||
|
||||
self._hash_db = HashDatabase(self.db_path)
|
||||
self._hash_db.initialize()
|
||||
self._ioc_db = IOCDatabase(self.db_path)
|
||||
self._ioc_db.initialize()
|
||||
count = self._hash_db.count()
|
||||
logger.info("Signature DB loaded: %d hash signatures", count)
|
||||
except Exception as exc:
|
||||
logger.error("Failed to load signature DB: %s", exc)
|
||||
self._loaded = True
|
||||
Reference in New Issue
Block a user