193 lines
6.7 KiB
Python
193 lines
6.7 KiB
Python
"""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
|