Cobra
ffd384f64b
Replace all sensor.* logger namespaces with __name__ (generic module identifiers instead of discoverable 'sensor.*' prefixes). Change hardcoded 'bb' API user to 'admin' in config and code defaults. Change hardcoded relay_user 'bb' to 'operator' — prevents network profiling from exposing tool identity via SSH config. Fixes #457, #458, #459
382 lines
14 KiB
Python
382 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Network relationship mapper — passive communication graph builder.
|
|
|
|
Builds a graph of all observed host-to-host communication with protocol,
|
|
byte count, and packet count metadata. Identifies roles (servers, clients,
|
|
admin workstations, printers) and generates Graphviz DOT output. Periodic
|
|
snapshots for change_detector consumption.
|
|
"""
|
|
|
|
import logging
|
|
import os
|
|
import socket
|
|
import sqlite3
|
|
import struct
|
|
import threading
|
|
import time
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from modules.base import BaseModule
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
TCP_PROTO = 6
|
|
UDP_PROTO = 17
|
|
|
|
# Role detection thresholds
|
|
SERVER_INBOUND_THRESHOLD = 5 # Unique source IPs connecting in
|
|
CLIENT_OUTBOUND_THRESHOLD = 10 # Unique dest IPs connected to
|
|
ADMIN_SSH_RDP_THRESHOLD = 3 # Unique SSH/RDP destinations
|
|
PRINTER_PORT = 9100
|
|
|
|
_DB_INIT = """
|
|
CREATE TABLE IF NOT EXISTS connections (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
src_ip TEXT NOT NULL,
|
|
dst_ip TEXT NOT NULL,
|
|
protocol TEXT NOT NULL,
|
|
port INTEGER NOT NULL,
|
|
bytes_total INTEGER DEFAULT 0,
|
|
packets INTEGER DEFAULT 0,
|
|
first_seen REAL NOT NULL,
|
|
last_seen REAL NOT NULL,
|
|
relationship_type TEXT DEFAULT '',
|
|
UNIQUE(src_ip, dst_ip, protocol, port)
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_conn_src ON connections(src_ip);
|
|
CREATE INDEX IF NOT EXISTS idx_conn_dst ON connections(dst_ip);
|
|
"""
|
|
|
|
FLUSH_INTERVAL = 45
|
|
SNAPSHOT_INTERVAL = 300 # 5 minutes
|
|
|
|
|
|
class NetworkMapper(BaseModule):
|
|
"""Build communication graph from all observed traffic."""
|
|
|
|
name = "network_mapper"
|
|
module_type = "passive"
|
|
priority = 150
|
|
requires_root = True
|
|
requires_capture_bus = True
|
|
|
|
def __init__(self, bus, state, config, engine=None):
|
|
super().__init__(bus, state, config, engine)
|
|
self._capture_bus = None
|
|
self._sub_queue = None
|
|
self._read_thread: Optional[threading.Thread] = None
|
|
self._flush_thread: Optional[threading.Thread] = None
|
|
self._snapshot_thread: Optional[threading.Thread] = None
|
|
self._db_path = self._resolve_db_path()
|
|
self._db_conn: Optional[sqlite3.Connection] = None
|
|
self._lock = threading.Lock()
|
|
# In-memory graph: (src_ip, dst_ip, proto, port) -> {bytes, packets, first, last}
|
|
self._graph: dict[tuple, dict] = {}
|
|
# Role tracking: ip -> set of connected IPs per direction
|
|
self._inbound: dict[str, set] = defaultdict(set) # dst -> set(src)
|
|
self._outbound: dict[str, set] = defaultdict(set) # src -> set(dst)
|
|
self._ssh_rdp_dests: dict[str, set] = defaultdict(set) # src -> set(dst) for SSH/RDP
|
|
self._printer_servers: set = set()
|
|
self._stats = {
|
|
"packets_processed": 0,
|
|
"unique_connections": 0,
|
|
"hosts_seen": 0,
|
|
"snapshots_taken": 0,
|
|
}
|
|
|
|
# ------------------------------------------------------------------
|
|
# BaseModule interface
|
|
# ------------------------------------------------------------------
|
|
|
|
def start(self) -> None:
|
|
if self._running:
|
|
return
|
|
|
|
self._init_db()
|
|
self._capture_bus = self.config.get("_capture_bus")
|
|
if self._capture_bus is None:
|
|
logger.error("NetworkMapper requires _capture_bus in config")
|
|
return
|
|
|
|
# Subscribe to all traffic
|
|
self._sub_queue = self._capture_bus.subscribe(
|
|
name=self.name, bpf_filter="", queue_depth=5000
|
|
)
|
|
|
|
self._running = True
|
|
self._start_time = time.time()
|
|
self._pid = os.getpid()
|
|
|
|
self._read_thread = threading.Thread(
|
|
target=self._read_loop, daemon=True, name="sensor-netmap-read"
|
|
)
|
|
self._read_thread.start()
|
|
|
|
self._flush_thread = threading.Thread(
|
|
target=self._flush_loop, daemon=True, name="sensor-netmap-flush"
|
|
)
|
|
self._flush_thread.start()
|
|
|
|
self._snapshot_thread = threading.Thread(
|
|
target=self._snapshot_loop, daemon=True, name="sensor-netmap-snapshot"
|
|
)
|
|
self._snapshot_thread.start()
|
|
|
|
self.state.set_module_status(self.name, "running", pid=self._pid)
|
|
self.bus.emit("MODULE_STARTED", {"module": self.name}, source_module=self.name)
|
|
logger.info("NetworkMapper started")
|
|
|
|
def stop(self) -> None:
|
|
if not self._running:
|
|
return
|
|
self._running = False
|
|
if self._capture_bus:
|
|
self._capture_bus.unsubscribe(self.name)
|
|
for t in (self._read_thread, self._flush_thread, self._snapshot_thread):
|
|
if t and t.is_alive():
|
|
t.join(timeout=5.0)
|
|
self._flush_to_db()
|
|
if self._db_conn:
|
|
self._db_conn.close()
|
|
self._db_conn = None
|
|
self.state.set_module_status(self.name, "stopped")
|
|
self.bus.emit("MODULE_STOPPED", {"module": self.name}, source_module=self.name)
|
|
logger.info("NetworkMapper stopped — stats: %s", self._stats)
|
|
|
|
def status(self) -> dict:
|
|
with self._lock:
|
|
hosts = set()
|
|
for (src, dst, _, _) in self._graph:
|
|
hosts.add(src)
|
|
hosts.add(dst)
|
|
self._stats["hosts_seen"] = len(hosts)
|
|
self._stats["unique_connections"] = len(self._graph)
|
|
return {
|
|
"running": self._running,
|
|
"pid": self._pid,
|
|
"uptime": time.time() - self._start_time if self._start_time else 0,
|
|
**self._stats,
|
|
}
|
|
|
|
def configure(self, config: dict) -> None:
|
|
self.config.update(config)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Packet processing
|
|
# ------------------------------------------------------------------
|
|
|
|
def _read_loop(self) -> None:
|
|
while self._running:
|
|
result = self._sub_queue.get(timeout=1.0)
|
|
if result is None:
|
|
continue
|
|
ts, pkt = result
|
|
self._stats["packets_processed"] += 1
|
|
try:
|
|
self._process_packet(pkt, ts)
|
|
except Exception:
|
|
pass
|
|
|
|
def _process_packet(self, pkt: bytes, ts: float) -> None:
|
|
"""Extract IP flow tuple and update graph."""
|
|
if len(pkt) < 34:
|
|
return
|
|
|
|
ethertype = struct.unpack("!H", pkt[12:14])[0]
|
|
|
|
# Handle 802.1Q
|
|
eth_offset = 14
|
|
if ethertype == 0x8100:
|
|
if len(pkt) < 38:
|
|
return
|
|
ethertype = struct.unpack("!H", pkt[16:18])[0]
|
|
eth_offset = 18
|
|
|
|
if ethertype != 0x0800: # IPv4 only
|
|
return
|
|
|
|
if len(pkt) < eth_offset + 20:
|
|
return
|
|
|
|
ip_header = pkt[eth_offset:]
|
|
ihl = (ip_header[0] & 0x0F) * 4
|
|
total_len = struct.unpack("!H", ip_header[2:4])[0]
|
|
ip_proto = ip_header[9]
|
|
src_ip = socket.inet_ntoa(ip_header[12:16])
|
|
dst_ip = socket.inet_ntoa(ip_header[16:20])
|
|
|
|
port = 0
|
|
proto_name = "other"
|
|
|
|
if ip_proto == TCP_PROTO and len(ip_header) >= ihl + 4:
|
|
_, dst_port = struct.unpack("!HH", ip_header[ihl:ihl + 4])
|
|
port = dst_port
|
|
proto_name = "tcp"
|
|
elif ip_proto == UDP_PROTO and len(ip_header) >= ihl + 4:
|
|
_, dst_port = struct.unpack("!HH", ip_header[ihl:ihl + 4])
|
|
port = dst_port
|
|
proto_name = "udp"
|
|
elif ip_proto == 1:
|
|
proto_name = "icmp"
|
|
|
|
pkt_bytes = total_len
|
|
|
|
key = (src_ip, dst_ip, proto_name, port)
|
|
with self._lock:
|
|
if key not in self._graph:
|
|
self._graph[key] = {
|
|
"bytes": 0, "packets": 0,
|
|
"first_seen": ts, "last_seen": ts,
|
|
}
|
|
entry = self._graph[key]
|
|
entry["bytes"] += pkt_bytes
|
|
entry["packets"] += 1
|
|
entry["last_seen"] = ts
|
|
|
|
# Track roles
|
|
self._inbound[dst_ip].add(src_ip)
|
|
self._outbound[src_ip].add(dst_ip)
|
|
if port in (22, 3389):
|
|
self._ssh_rdp_dests[src_ip].add(dst_ip)
|
|
if port == PRINTER_PORT:
|
|
self._printer_servers.add(dst_ip)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Role classification
|
|
# ------------------------------------------------------------------
|
|
|
|
def _classify_role(self, ip: str) -> str:
|
|
"""Classify a host's role based on observed traffic patterns."""
|
|
roles = []
|
|
if len(self._inbound.get(ip, set())) >= SERVER_INBOUND_THRESHOLD:
|
|
roles.append("server")
|
|
if len(self._ssh_rdp_dests.get(ip, set())) >= ADMIN_SSH_RDP_THRESHOLD:
|
|
roles.append("admin_workstation")
|
|
if ip in self._printer_servers:
|
|
roles.append("printer")
|
|
if len(self._outbound.get(ip, set())) >= CLIENT_OUTBOUND_THRESHOLD:
|
|
roles.append("client")
|
|
return ",".join(roles) if roles else "unknown"
|
|
|
|
# ------------------------------------------------------------------
|
|
# Graphviz output
|
|
# ------------------------------------------------------------------
|
|
|
|
def generate_dot(self) -> str:
|
|
"""Generate Graphviz DOT format of the communication graph."""
|
|
lines = ['digraph network {', ' rankdir=LR;', ' node [shape=box];']
|
|
|
|
with self._lock:
|
|
hosts = set()
|
|
for (src, dst, proto, port) in self._graph:
|
|
hosts.add(src)
|
|
hosts.add(dst)
|
|
|
|
# Node declarations with role colors
|
|
role_colors = {
|
|
"server": "lightblue",
|
|
"admin_workstation": "orange",
|
|
"printer": "lightgreen",
|
|
"client": "lightyellow",
|
|
}
|
|
for ip in sorted(hosts):
|
|
role = self._classify_role(ip)
|
|
color = "white"
|
|
for r, c in role_colors.items():
|
|
if r in role:
|
|
color = c
|
|
break
|
|
label = f"{ip}\\n[{role}]"
|
|
lines.append(f' "{ip}" [label="{label}", style=filled, fillcolor={color}];')
|
|
|
|
# Edges
|
|
for (src, dst, proto, port), data in self._graph.items():
|
|
label = f"{proto}/{port}\\n{data['packets']}pkts"
|
|
lines.append(f' "{src}" -> "{dst}" [label="{label}"];')
|
|
|
|
lines.append('}')
|
|
return '\n'.join(lines)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Database
|
|
# ------------------------------------------------------------------
|
|
|
|
def _resolve_db_path(self) -> str:
|
|
base = self.config.get("db_dir", os.path.join(
|
|
os.path.expanduser("~"), ".implant"
|
|
))
|
|
return os.path.join(base, "network_mapper.db")
|
|
|
|
def _init_db(self) -> None:
|
|
Path(self._db_path).parent.mkdir(parents=True, exist_ok=True)
|
|
self._db_conn = sqlite3.connect(self._db_path, check_same_thread=False)
|
|
self._db_conn.execute("PRAGMA journal_mode=WAL")
|
|
self._db_conn.execute("PRAGMA synchronous=NORMAL")
|
|
self._db_conn.executescript(_DB_INIT)
|
|
self._db_conn.commit()
|
|
|
|
def _flush_to_db(self) -> None:
|
|
"""Flush in-memory graph to SQLite."""
|
|
with self._lock:
|
|
snapshot = dict(self._graph)
|
|
|
|
if not snapshot or not self._db_conn:
|
|
return
|
|
|
|
try:
|
|
with self._db_conn:
|
|
for (src, dst, proto, port), data in snapshot.items():
|
|
rel = self._classify_role(dst)
|
|
self._db_conn.execute(
|
|
"""INSERT INTO connections
|
|
(src_ip, dst_ip, protocol, port, bytes_total, packets,
|
|
first_seen, last_seen, relationship_type)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
ON CONFLICT(src_ip, dst_ip, protocol, port)
|
|
DO UPDATE SET
|
|
bytes_total = bytes_total + excluded.bytes_total,
|
|
packets = packets + excluded.packets,
|
|
last_seen = MAX(last_seen, excluded.last_seen)""",
|
|
(src, dst, proto, port, data["bytes"], data["packets"],
|
|
data["first_seen"], data["last_seen"], rel),
|
|
)
|
|
except Exception:
|
|
logger.exception("Failed to flush connections to DB")
|
|
|
|
def _flush_loop(self) -> None:
|
|
while self._running:
|
|
time.sleep(FLUSH_INTERVAL)
|
|
try:
|
|
self._flush_to_db()
|
|
except Exception:
|
|
logger.exception("Flush loop error")
|
|
|
|
def _snapshot_loop(self) -> None:
|
|
"""Periodic snapshots for change_detector consumption."""
|
|
while self._running:
|
|
time.sleep(SNAPSHOT_INTERVAL)
|
|
try:
|
|
dot_output = self.generate_dot()
|
|
snapshot_dir = os.path.join(
|
|
os.path.dirname(self._db_path), "snapshots"
|
|
)
|
|
Path(snapshot_dir).mkdir(parents=True, exist_ok=True)
|
|
filename = f"netmap_{int(time.time())}.dot"
|
|
filepath = os.path.join(snapshot_dir, filename)
|
|
with open(filepath, "w") as f:
|
|
f.write(dot_output)
|
|
self._stats["snapshots_taken"] += 1
|
|
|
|
# Publish snapshot event for change_detector
|
|
self.bus.emit("CHANGE_DETECTED", {
|
|
"module": self.name,
|
|
"snapshot_path": filepath,
|
|
"hosts_count": self._stats["hosts_seen"],
|
|
"connections_count": self._stats["unique_connections"],
|
|
}, source_module=self.name)
|
|
|
|
except Exception:
|
|
logger.exception("Snapshot loop error")
|