remove infra.md.example, infra.md is the source of truth
This commit is contained in:
Azreen Jamal
259
ayn-antivirus/ayn_antivirus/signatures/db/ioc_db.py
Normal file
259
ayn-antivirus/ayn_antivirus/signatures/db/ioc_db.py
Normal file
@@ -0,0 +1,259 @@
|
||||
"""SQLite-backed Indicator of Compromise (IOC) database for AYN Antivirus.
|
||||
|
||||
Stores malicious IPs, domains, and URLs sourced from threat-intelligence
|
||||
feeds so that the network scanner and detectors can perform real-time
|
||||
lookups.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import sqlite3
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional, Sequence, Set, Tuple
|
||||
|
||||
from ayn_antivirus.constants import DEFAULT_DB_PATH
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Schema
|
||||
# ---------------------------------------------------------------------------
|
||||
_SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS ioc_ips (
|
||||
ip TEXT PRIMARY KEY,
|
||||
threat_name TEXT NOT NULL DEFAULT '',
|
||||
type TEXT NOT NULL DEFAULT 'C2',
|
||||
source TEXT NOT NULL DEFAULT '',
|
||||
added_date TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_ioc_ips_source ON ioc_ips(source);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS ioc_domains (
|
||||
domain TEXT PRIMARY KEY,
|
||||
threat_name TEXT NOT NULL DEFAULT '',
|
||||
type TEXT NOT NULL DEFAULT 'C2',
|
||||
source TEXT NOT NULL DEFAULT '',
|
||||
added_date TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_ioc_domains_source ON ioc_domains(source);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS ioc_urls (
|
||||
url TEXT PRIMARY KEY,
|
||||
threat_name TEXT NOT NULL DEFAULT '',
|
||||
type TEXT NOT NULL DEFAULT 'malware_distribution',
|
||||
source TEXT NOT NULL DEFAULT '',
|
||||
added_date TEXT NOT NULL DEFAULT (datetime('now'))
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_ioc_urls_source ON ioc_urls(source);
|
||||
"""
|
||||
|
||||
|
||||
class IOCDatabase:
|
||||
"""Manage a local SQLite store of Indicators of Compromise.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
db_path:
|
||||
Path to the SQLite file. Shares the same file as
|
||||
:class:`HashDatabase` by default; each uses its own tables.
|
||||
"""
|
||||
|
||||
_VALID_TABLES: frozenset = frozenset({"ioc_ips", "ioc_domains", "ioc_urls"})
|
||||
|
||||
def __init__(self, db_path: str | Path = DEFAULT_DB_PATH) -> None:
|
||||
self.db_path = Path(db_path)
|
||||
self._conn: Optional[sqlite3.Connection] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def initialize(self) -> None:
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._conn = sqlite3.connect(str(self.db_path), check_same_thread=False)
|
||||
self._conn.row_factory = sqlite3.Row
|
||||
self._conn.execute("PRAGMA journal_mode=WAL")
|
||||
self._conn.executescript(_SCHEMA)
|
||||
self._conn.commit()
|
||||
logger.info(
|
||||
"IOCDatabase opened: %s (IPs=%d, domains=%d, URLs=%d)",
|
||||
self.db_path,
|
||||
self._count("ioc_ips"),
|
||||
self._count("ioc_domains"),
|
||||
self._count("ioc_urls"),
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
if self._conn:
|
||||
self._conn.close()
|
||||
self._conn = None
|
||||
|
||||
@property
|
||||
def conn(self) -> sqlite3.Connection:
|
||||
if self._conn is None:
|
||||
self.initialize()
|
||||
assert self._conn is not None
|
||||
return self._conn
|
||||
|
||||
def _count(self, table: str) -> int:
|
||||
if table not in self._VALID_TABLES:
|
||||
raise ValueError(f"Invalid table name: {table}")
|
||||
return self.conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0]
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# IPs
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def add_ip(
|
||||
self,
|
||||
ip: str,
|
||||
threat_name: str = "",
|
||||
type: str = "C2",
|
||||
source: str = "",
|
||||
) -> None:
|
||||
self.conn.execute(
|
||||
"INSERT OR REPLACE INTO ioc_ips (ip, threat_name, type, source, added_date) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
(ip, threat_name, type, source, datetime.utcnow().isoformat()),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def bulk_add_ips(
|
||||
self,
|
||||
records: Sequence[Tuple[str, str, str, str]],
|
||||
) -> int:
|
||||
"""Bulk-insert IPs. Each tuple: ``(ip, threat_name, type, source)``.
|
||||
|
||||
Returns the number of **new** rows actually inserted.
|
||||
"""
|
||||
if not records:
|
||||
return 0
|
||||
now = datetime.utcnow().isoformat()
|
||||
rows = [(ip, tn, t, src, now) for ip, tn, t, src in records]
|
||||
before = self._count("ioc_ips")
|
||||
self.conn.executemany(
|
||||
"INSERT OR IGNORE INTO ioc_ips (ip, threat_name, type, source, added_date) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
rows,
|
||||
)
|
||||
self.conn.commit()
|
||||
return self._count("ioc_ips") - before
|
||||
|
||||
def lookup_ip(self, ip: str) -> Optional[Dict[str, Any]]:
|
||||
row = self.conn.execute(
|
||||
"SELECT * FROM ioc_ips WHERE ip = ?", (ip,)
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def get_all_malicious_ips(self) -> Set[str]:
|
||||
"""Return every stored malicious IP as a set for fast membership tests."""
|
||||
rows = self.conn.execute("SELECT ip FROM ioc_ips").fetchall()
|
||||
return {row[0] for row in rows}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Domains
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def add_domain(
|
||||
self,
|
||||
domain: str,
|
||||
threat_name: str = "",
|
||||
type: str = "C2",
|
||||
source: str = "",
|
||||
) -> None:
|
||||
self.conn.execute(
|
||||
"INSERT OR REPLACE INTO ioc_domains (domain, threat_name, type, source, added_date) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
(domain.lower(), threat_name, type, source, datetime.utcnow().isoformat()),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def bulk_add_domains(
|
||||
self,
|
||||
records: Sequence[Tuple[str, str, str, str]],
|
||||
) -> int:
|
||||
"""Bulk-insert domains. Each tuple: ``(domain, threat_name, type, source)``.
|
||||
|
||||
Returns the number of **new** rows actually inserted.
|
||||
"""
|
||||
if not records:
|
||||
return 0
|
||||
now = datetime.utcnow().isoformat()
|
||||
rows = [(d.lower(), tn, t, src, now) for d, tn, t, src in records]
|
||||
before = self._count("ioc_domains")
|
||||
self.conn.executemany(
|
||||
"INSERT OR IGNORE INTO ioc_domains (domain, threat_name, type, source, added_date) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
rows,
|
||||
)
|
||||
self.conn.commit()
|
||||
return self._count("ioc_domains") - before
|
||||
|
||||
def lookup_domain(self, domain: str) -> Optional[Dict[str, Any]]:
|
||||
row = self.conn.execute(
|
||||
"SELECT * FROM ioc_domains WHERE domain = ?", (domain.lower(),)
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def get_all_malicious_domains(self) -> Set[str]:
|
||||
"""Return every stored malicious domain as a set."""
|
||||
rows = self.conn.execute("SELECT domain FROM ioc_domains").fetchall()
|
||||
return {row[0] for row in rows}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# URLs
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def add_url(
|
||||
self,
|
||||
url: str,
|
||||
threat_name: str = "",
|
||||
type: str = "malware_distribution",
|
||||
source: str = "",
|
||||
) -> None:
|
||||
self.conn.execute(
|
||||
"INSERT OR REPLACE INTO ioc_urls (url, threat_name, type, source, added_date) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
(url, threat_name, type, source, datetime.utcnow().isoformat()),
|
||||
)
|
||||
self.conn.commit()
|
||||
|
||||
def bulk_add_urls(
|
||||
self,
|
||||
records: Sequence[Tuple[str, str, str, str]],
|
||||
) -> int:
|
||||
"""Bulk-insert URLs. Each tuple: ``(url, threat_name, type, source)``.
|
||||
|
||||
Returns the number of **new** rows actually inserted.
|
||||
"""
|
||||
if not records:
|
||||
return 0
|
||||
now = datetime.utcnow().isoformat()
|
||||
rows = [(u, tn, t, src, now) for u, tn, t, src in records]
|
||||
before = self._count("ioc_urls")
|
||||
self.conn.executemany(
|
||||
"INSERT OR IGNORE INTO ioc_urls (url, threat_name, type, source, added_date) "
|
||||
"VALUES (?, ?, ?, ?, ?)",
|
||||
rows,
|
||||
)
|
||||
self.conn.commit()
|
||||
return self._count("ioc_urls") - before
|
||||
|
||||
def lookup_url(self, url: str) -> Optional[Dict[str, Any]]:
|
||||
row = self.conn.execute(
|
||||
"SELECT * FROM ioc_urls WHERE url = ?", (url,)
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Aggregate stats
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_stats(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"ips": self._count("ioc_ips"),
|
||||
"domains": self._count("ioc_domains"),
|
||||
"urls": self._count("ioc_urls"),
|
||||
}
|
||||
Reference in New Issue
Block a user