defense.py

from __future__ import annotations
from typing import List, Tuple, Optional
import sqlite3
import json
import time
import threading
import signal
import os
import hashlib
import psutil
import numpy as np
from sklearn.ensemble import IsolationForest
import joblib


class DatabaseHelper:
	"""Handles SQLite operations for feature storage, pruning and quarantine."""
	def __init__(self, db_path: str) -> None:
		self.db_path = db_path
		self._ensure_db()

	def _ensure_db(self) -> None:
		with sqlite3.connect(self.db_path) as con:
			cur = con.cursor()
			# Main feature table
			cur.execute("""
			CREATE TABLE IF NOT EXISTS features (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				ts INTEGER,
				pid INTEGER,
				ppid INTEGER,
				username TEXT,
				name TEXT,
				cmdline TEXT,
				sha256_cmdline TEXT,
				cwd TEXT,
				cpu REAL,
				mem REAL,
				io_read INTEGER,
				io_write INTEGER,
				num_threads INTEGER,
				num_fds INTEGER,
				num_conns INTEGER,
				num_envvars INTEGER,
				num_modules INTEGER
			)""")
			# Labels and quarantines tables
			cur.execute("""
			CREATE TABLE IF NOT EXISTS labels (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				feature_id INTEGER,
				label INTEGER,
				source TEXT,
				ts INTEGER
			)""")
			cur.execute("""
			CREATE TABLE IF NOT EXISTS quarantines (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				ts INTEGER,
				pid INTEGER,
				snapshot_json TEXT
			)""")
			# Environment variables per feature
			cur.execute("""
			CREATE TABLE IF NOT EXISTS feature_envvars (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				feature_id INTEGER,
				key TEXT,
				value TEXT
			)""")
			# Loaded modules per feature
			cur.execute("""
			CREATE TABLE IF NOT EXISTS feature_modules (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				feature_id INTEGER,
				module_name TEXT,
				module_path TEXT
			)""")
			con.commit()

	def insert_feature(self, feature_row: Tuple, env_vars: List[Tuple[str, str]], modules: List[Tuple[str, str]]) -> None:
		"""Insert main feature row + child table rows."""
		with sqlite3.connect(self.db_path) as con:
			cur = con.cursor()
			cur.execute("""
			INSERT INTO features (ts,pid,ppid,username,name,cmdline,sha256_cmdline,cwd,cpu,mem,io_read,io_write,num_threads,num_fds,num_conns,num_envvars,num_modules)
			VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
			""", feature_row)
			feature_id = cur.lastrowid
			for k, v in env_vars:
				cur.execute("INSERT INTO feature_envvars (feature_id,key,value) VALUES (?,?,?)", (feature_id, k, v))
			for m_name, m_path in modules:
				cur.execute("INSERT INTO feature_modules (feature_id,module_name,module_path) VALUES (?,?,?)", (feature_id, m_name, m_path))
			con.commit()

	def prune_if_needed(self, max_rows: int) -> None:
		with sqlite3.connect(self.db_path) as con:
			cur = con.cursor()
			cur.execute("SELECT COUNT(*) FROM features")
			(count,) = cur.fetchone()
			if count > max_rows:
				to_delete = count - max_rows
				cur.execute("DELETE FROM features WHERE id IN (SELECT id FROM features ORDER BY ts ASC LIMIT ?)", (to_delete,))
				con.commit()

	def insert_quarantine(self, pid: int, snapshot: dict) -> None:
		with sqlite3.connect(self.db_path) as con:
			cur = con.cursor()
			cur.execute("INSERT INTO quarantines (ts,pid,snapshot_json) VALUES (?,?,?)",
						(int(time.time()), pid, json.dumps(snapshot)))
			con.commit()


class AdaptiveDefender:
	"""Class-based adaptive malware defense with monitoring, detection, and retraining."""
	def __init__(self, db_path: str, model_path: str) -> None:
		self.db = DatabaseHelper(db_path)
		self.model_path = model_path
		self.model_lock = threading.Lock()
		self.stop_event = threading.Event()
		self.model: Optional[IsolationForest] = None
		self.load_or_init_model()

	def now_ts(self) -> int:
		return int(time.time())

	def load_or_init_model(self) -> None:
		if os.path.exists(self.model_path):
			self.model = joblib.load(self.model_path)
		else:
			dummy = IsolationForest(contamination=0.01, random_state=0)
			dummy.fit(np.zeros((10,9)))
			joblib.dump(dummy, self.model_path)
			self.model = dummy

	def extract_features(self, proc: psutil.Process) -> Optional[Tuple[Tuple, List[Tuple[str,str]], List[Tuple[str,str]]]]:
		try:
			with proc.oneshot():
				cpu = proc.cpu_percent(interval=0.1)
				mem = proc.memory_percent()
				io_read = io_write = 0
				try:
					io = proc.io_counters()
					io_read = io.read_bytes if io else 0
					io_write = io.write_bytes if io else 0
				except Exception:
					pass

				threads = proc.num_threads()
				num_fds = getattr(proc, "num_fds", lambda: 0)()
				num_conns = len(proc.net_connections(kind='inet'))

				cmdline_str = ' '.join(proc.cmdline()) if proc.is_running() else ''
				sha_cmd = hashlib.sha256(cmdline_str.encode()).hexdigest()
				username = proc.username()
				name = proc.name()[:200]
				cwd = proc.cwd() if proc.is_running() else ''

				try:
					env = proc.environ()
					env_list = list(env.items())
				except Exception:
					env_list = []
				num_env = len(env_list)

				mod_list = []
				try:
					for m in proc.memory_maps():
						mod_list.append((m.path.split('/')[-1], m.path))
				except Exception:
					pass
				num_mod = len(mod_list)

				feature_row = (
					self.now_ts(), proc.pid, proc.ppid(), username, name,
					cmdline_str, sha_cmd, cwd,
					float(cpu), float(mem), int(io_read), int(io_write),
					int(threads), int(num_fds), int(num_conns),
					num_env, num_mod
				)
				return feature_row, env_list, mod_list
		except Exception:
			return None

	def contain_process(proc, reason='unknown'):
		"""
		Safely gather process info and optionally quarantine it.
		Works on Linux, Windows, macOS with fallbacks.
		"""
		if not proc.is_running():
			return
	
		# --- Gather modules / memory mappings ---
		modules = []
		try:
			# Primary: memory_maps() (works on Linux, some Unix)
			modules = [(m.path.split('/')[-1], m.path) for m in proc.memory_maps()]
		except (AttributeError, psutil.AccessDenied):
			# Fallback: use open_files() if memory_maps unavailable
			try:
				modules = [(f.path.split('/')[-1], f.path) for f in proc.open_files()]
			except (psutil.AccessDenied, Exception):
				modules = []
	
		# --- Gather environment variables ---
		envvars = {}
		try:
			envvars = proc.environ()
		except (psutil.AccessDenied, AttributeError):
			envvars = {}
	
		# --- Gather basic stats ---
		try:
			info = {
				'pid': proc.pid,
				'name': proc.name(),
				'username': proc.username(),
				'cmdline': proc.cmdline(),
				'cwd': proc.cwd() if hasattr(proc, 'cwd') else '',
				'cpu': proc.cpu_percent(interval=0.1),
				'mem': proc.memory_percent(),
				'io_read': getattr(proc, 'io_counters', lambda: None)().read_bytes if hasattr(proc, 'io_counters') else 0,
				'io_write': getattr(proc, 'io_counters', lambda: None)().write_bytes if hasattr(proc, 'io_counters') else 0,
				'num_threads': proc.num_threads(),
				'num_fds': proc.num_fds() if hasattr(proc, 'num_fds') else 0,
				'num_conns': len(proc.connections()) if hasattr(proc, 'connections') else 0,
				'num_envvars': len(envvars),
				'modules': modules,
				'reason': reason
			}
		except (psutil.NoSuchProcess, psutil.AccessDenied):
			return  # Process died or inaccessible
	
		# --- Optionally: quarantine or log process ---
		# Example: save to DB or kill
		# db_save(info)
		# proc.terminate()
	
		return info

	def monitor_loop(self, poll_interval: float = 1.0) -> None:
		while not self.stop_event.is_set():
			start = time.time()
			for proc in psutil.process_iter(['pid', 'name']):
				data = self.extract_features(proc)
				if data:
					feature_row, env_list, mod_list = data
					self.db.insert_feature(feature_row, env_list, mod_list)
			self.db.prune_if_needed(200000)
			time.sleep(max(0, poll_interval - (time.time() - start)))

	def detector_loop(self, scan_interval: float = 2.0) -> None:
		while not self.stop_event.is_set():
			if self.model is None:
				time.sleep(scan_interval)
				continue

			records = []
			for proc in psutil.process_iter(['pid', 'name']):
				data = self.extract_features(proc)
				if not data:
					continue
				feat_row, _, _ = data
				vec = np.array(feat_row[8:]).astype(float)  # take the 9-dimensional vector
				records.append((proc, vec))

			if records:
				with self.model_lock:
					X = np.vstack([r[1] for r in records])
					preds = self.model.predict(X)
					scores = self.model.decision_function(X)

				for (proc, _), pred, score in zip(records, preds, scores):
					if pred == -1:
						self.contain_process(proc, f'anomaly_score={score}')

			time.sleep(scan_interval)

	def fetch_training_data(self, window_minutes: int = 360) -> np.ndarray:
		cutoff = self.now_ts() - int(window_minutes * 60)
		with sqlite3.connect(self.db.db_path) as con:
			cur = con.cursor()
			cur.execute("""
				SELECT cpu, mem, io_read, io_write, num_threads, num_fds, num_conns, num_envvars, num_modules
				FROM features WHERE ts >= ? ORDER BY ts DESC
			""", (cutoff,))
			rows = cur.fetchall()
		return np.array(rows) if rows else np.empty((0, 9), dtype=float)

	def retrain_worker(self, interval_sec: int = 300, min_samples: int = 500) -> None:
		while not self.stop_event.is_set():
			with self.model_lock:
				X = self.fetch_training_data()
				if X.shape[0] >= min_samples:
					X_train = np.vstack([X, self.augment_features(X, 0.1)])
					model = IsolationForest(contamination=0.01, random_state=42)
					model.fit(X_train)
					joblib.dump(model, self.model_path)
					self.model = model

			for _ in range(int(interval_sec / 2)):
				if self.stop_event.is_set():
					break
				time.sleep(2)

	def augment_features(self, X: np.ndarray, factor: float = 0.05) -> np.ndarray:
		if X.size == 0:
			return X
		n, d = X.shape
		num_aug = max(1, int(n * factor))
		idx = np.random.randint(0, n, size=num_aug)
		samples = X[idx] + 0.01 * np.random.randn(num_aug, d) * np.maximum(1.0, np.abs(X[idx]))
		return np.clip(samples, a_min=0, a_max=None)

	def shutdown(self) -> None:
		self.stop_event.set()


def main() -> None:
	defender = AdaptiveDefender('behavior.db', 'behavior_model.joblib')

	mon_thread = threading.Thread(target=defender.monitor_loop, daemon=True)
	det_thread = threading.Thread(target=defender.detector_loop, daemon=True)
	train_thread = threading.Thread(target=defender.retrain_worker, daemon=True)

	mon_thread.start()
	det_thread.start()
	train_thread.start()

	def handle_signal(signum, frame):
		defender.shutdown()

	signal.signal(signal.SIGINT, handle_signal)
	signal.signal(signal.SIGTERM, handle_signal)

	try:
		while not defender.stop_event.is_set():
			time.sleep(1)
	except KeyboardInterrupt:
		defender.shutdown()


if __name__ == '__main__':
	main()