diff options
32 files changed, 713 insertions, 557 deletions
diff --git a/generalresearch/managers/network/label.py b/generalresearch/managers/network/label.py index 0405716..65c63e5 100644 --- a/generalresearch/managers/network/label.py +++ b/generalresearch/managers/network/label.py @@ -1,6 +1,7 @@ from datetime import datetime, timezone, timedelta from typing import Collection, Optional, List +from psycopg import sql from pydantic import TypeAdapter, IPvAnyNetwork from generalresearch.managers.base import PostgresManager @@ -14,7 +15,8 @@ from generalresearch.models.network.label import IPLabel, IPLabelKind, IPLabelSo class IPLabelManager(PostgresManager): def create(self, ip_label: IPLabel) -> IPLabel: - query = """ + query = sql.SQL( + """ INSERT INTO network_iplabel ( ip, labeled_at, created_at, label_kind, source, confidence, @@ -24,6 +26,7 @@ class IPLabelManager(PostgresManager): %(label_kind)s, %(source)s, %(confidence)s, %(provider)s, %(metadata)s ) RETURNING id;""" + ) params = ip_label.model_dump_postgres() with self.pg_config.make_connection() as conn: with conn.cursor() as c: diff --git a/generalresearch/managers/network/mtr.py b/generalresearch/managers/network/mtr.py index 496d3f0..35e4871 100644 --- a/generalresearch/managers/network/mtr.py +++ b/generalresearch/managers/network/mtr.py @@ -1,18 +1,19 @@ from typing import Optional -from psycopg import Cursor +from psycopg import Cursor, sql from generalresearch.managers.base import PostgresManager -from generalresearch.models.network.tool_run import MtrRun +from generalresearch.models.network.tool_run import MTRRun -class MtrManager(PostgresManager): +class MTRRunManager(PostgresManager): - def _create(self, run: MtrRun, c: Optional[Cursor] = None) -> None: + def _create(self, run: MTRRun, c: Optional[Cursor] = None) -> None: """ Do not use this directly. Must only be used in the context of a toolrun """ - query = """ + query = sql.SQL( + """ INSERT INTO network_mtr ( run_id, source_ip, facility_id, protocol, port, parsed @@ -22,9 +23,11 @@ class MtrManager(PostgresManager): %(protocol)s, %(port)s, %(parsed)s ); """ + ) params = run.model_dump_postgres() - query_hops = """ + query_hops = sql.SQL( + """ INSERT INTO network_mtrhop ( hop, ip, domain, asn, mtr_run_id ) VALUES ( @@ -32,6 +35,7 @@ class MtrManager(PostgresManager): %(asn)s, %(mtr_run_id)s ) """ + ) mtr_run = run.parsed params_hops = [h.model_dump_postgres(run_id=run.id) for h in mtr_run.hops] diff --git a/generalresearch/managers/network/nmap.py b/generalresearch/managers/network/nmap.py index 9cbc283..0995a32 100644 --- a/generalresearch/managers/network/nmap.py +++ b/generalresearch/managers/network/nmap.py @@ -1,19 +1,20 @@ from typing import Optional -from psycopg import Cursor +from psycopg import Cursor, sql from generalresearch.managers.base import PostgresManager -from generalresearch.models.network.tool_run import PortScanRun +from generalresearch.models.network.tool_run import NmapRun -class NmapManager(PostgresManager): +class NmapRunManager(PostgresManager): - def _create(self, run: PortScanRun, c: Optional[Cursor] = None) -> None: + def _create(self, run: NmapRun, c: Optional[Cursor] = None) -> None: """ - Insert a PortScan + PortScanPorts from a Pydantic NmapRun. + Insert a PortScan + PortScanPorts from a Pydantic NmapResult. Do not use this directly. Must only be used in the context of a toolrun """ - query = """ + query = sql.SQL( + """ INSERT INTO network_portscan ( run_id, xml_version, host_state, host_state_reason, latency_ms, distance, @@ -29,9 +30,11 @@ class NmapManager(PostgresManager): %(started_at)s, %(ip)s ); """ + ) params = run.model_dump_postgres() - query_ports = """ + query_ports = sql.SQL( + """ INSERT INTO network_portscanport ( port_scan_id, protocol, port, state, reason, reason_ttl, @@ -42,6 +45,7 @@ class NmapManager(PostgresManager): %(service_name)s ) """ + ) nmap_run = run.parsed params_ports = [p.model_dump_postgres(run_id=run.id) for p in nmap_run.ports] diff --git a/generalresearch/managers/network/rdns.py b/generalresearch/managers/network/rdns.py index 0b9b7b6..3543180 100644 --- a/generalresearch/managers/network/rdns.py +++ b/generalresearch/managers/network/rdns.py @@ -3,12 +3,12 @@ from typing import Optional from psycopg import Cursor from generalresearch.managers.base import PostgresManager -from generalresearch.models.network.tool_run import RDnsRun +from generalresearch.models.network.tool_run import RDNSRun -class RdnsManager(PostgresManager): +class RDNSRunManager(PostgresManager): - def _create(self, run: RDnsRun, c: Optional[Cursor] = None) -> None: + def _create(self, run: RDNSRun, c: Optional[Cursor] = None) -> None: """ Do not use this directly. Must only be used in the context of a toolrun """ @@ -23,4 +23,9 @@ class RdnsManager(PostgresManager): ); """ params = run.model_dump_postgres() - c.execute(query, params)
\ No newline at end of file + if c: + c.execute(query, params) + else: + with self.pg_config.make_connection() as conn: + with conn.cursor() as c: + c.execute(query, params) diff --git a/generalresearch/managers/network/tool_run.py b/generalresearch/managers/network/tool_run.py index 6280221..33853a0 100644 --- a/generalresearch/managers/network/tool_run.py +++ b/generalresearch/managers/network/tool_run.py @@ -3,16 +3,12 @@ from typing import Collection from psycopg import Cursor, sql from generalresearch.managers.base import PostgresManager, Permission -from generalresearch.models.network.rdns import RDNSResult -from generalresearch.models.network.tool_run import ( - ToolRun, - PortScanRun, - RDnsRun, - MtrRun, -) -from generalresearch.managers.network.nmap import NmapManager -from generalresearch.managers.network.rdns import RdnsManager -from generalresearch.managers.network.mtr import MtrManager + +from generalresearch.managers.network.nmap import NmapRunManager +from generalresearch.managers.network.rdns import RDNSRunManager +from generalresearch.managers.network.mtr import MTRRunManager +from generalresearch.models.network.rdns.result import RDNSResult +from generalresearch.models.network.tool_run import NmapRun, RDNSRun, MTRRun from generalresearch.pg_helper import PostgresConfig @@ -23,11 +19,11 @@ class ToolRunManager(PostgresManager): permissions: Collection[Permission] = None, ): super().__init__(pg_config=pg_config, permissions=permissions) - self.nmap_manager = NmapManager(self.pg_config) - self.rdns_manager = RdnsManager(self.pg_config) - self.mtr_manager = MtrManager(self.pg_config) + self.nmap_manager = NmapRunManager(self.pg_config) + self.rdns_manager = RDNSRunManager(self.pg_config) + self.mtr_manager = MTRRunManager(self.pg_config) - def create_tool_run(self, run: PortScanRun | RDnsRun | MtrRun, c: Cursor): + def create_tool_run(self, run: NmapRun | RDNSRun | MTRRun, c: Cursor): query = sql.SQL( """ INSERT INTO network_toolrun ( @@ -50,9 +46,9 @@ class ToolRunManager(PostgresManager): run.id = run_id return None - def create_portscan_run(self, run: PortScanRun) -> PortScanRun: + def create_nmap_run(self, run: NmapRun) -> NmapRun: """ - Insert a PortScan + PortScanPorts from a Pydantic NmapRun. + Insert a PortScan + PortScanPorts from a Pydantic NmapResult. """ with self.pg_config.make_connection() as conn: with conn.cursor() as c: @@ -60,7 +56,7 @@ class ToolRunManager(PostgresManager): self.nmap_manager._create(run, c=c) return run - def get_portscan_run(self, id: int) -> PortScanRun: + def get_nmap_run(self, id: int) -> NmapRun: query = """ SELECT tr.*, np.parsed FROM network_toolrun tr @@ -69,9 +65,9 @@ class ToolRunManager(PostgresManager): """ params = {"id": id} res = self.pg_config.execute_sql_query(query, params)[0] - return PortScanRun.model_validate(res) + return NmapRun.model_validate(res) - def create_rdns_run(self, run: RDnsRun) -> RDnsRun: + def create_rdns_run(self, run: RDNSRun) -> RDNSRun: """ Insert a RDnsRun + RDNSResult """ @@ -81,7 +77,7 @@ class ToolRunManager(PostgresManager): self.rdns_manager._create(run, c=c) return run - def get_rdns_run(self, id: int) -> RDnsRun: + def get_rdns_run(self, id: int) -> RDNSRun: query = """ SELECT tr.*, hostnames FROM network_toolrun tr @@ -94,16 +90,16 @@ class ToolRunManager(PostgresManager): {"ip": res["ip"], "hostnames": res["hostnames"]} ) res["parsed"] = parsed - return RDnsRun.model_validate(res) + return RDNSRun.model_validate(res) - def create_mtr_run(self, run: MtrRun) -> MtrRun: + def create_mtr_run(self, run: MTRRun) -> MTRRun: with self.pg_config.make_connection() as conn: with conn.cursor() as c: self.create_tool_run(run, c) self.mtr_manager._create(run, c=c) return run - def get_mtr_run(self, id: int) -> MtrRun: + def get_mtr_run(self, id: int) -> MTRRun: query = """ SELECT tr.*, mtr.parsed, mtr.source_ip, mtr.facility_id FROM network_toolrun tr @@ -112,4 +108,4 @@ class ToolRunManager(PostgresManager): """ params = {"id": id} res = self.pg_config.execute_sql_query(query, params)[0] - return MtrRun.model_validate(res) + return MTRRun.model_validate(res) diff --git a/generalresearch/models/network/mtr/__init__.py b/generalresearch/models/network/mtr/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/generalresearch/models/network/mtr/__init__.py diff --git a/generalresearch/models/network/mtr/command.py b/generalresearch/models/network/mtr/command.py new file mode 100644 index 0000000..e3ab903 --- /dev/null +++ b/generalresearch/models/network/mtr/command.py @@ -0,0 +1,75 @@ +import subprocess +from typing import List, Optional + +from generalresearch.models.network.definitions import IPProtocol +from generalresearch.models.network.mtr.parser import parse_mtr_output +from generalresearch.models.network.mtr.result import MTRResult + +SUPPORTED_PROTOCOLS = { + IPProtocol.TCP, + IPProtocol.UDP, + IPProtocol.SCTP, + IPProtocol.ICMP, +} +PROTOCOLS_W_PORT = {IPProtocol.TCP, IPProtocol.UDP, IPProtocol.SCTP} + + +def build_mtr_command( + ip: str, + protocol: Optional[IPProtocol] = None, + port: Optional[int] = None, + report_cycles: int = 10, +) -> str: + # https://manpages.ubuntu.com/manpages/focal/man8/mtr.8.html + # e.g. "mtr -r -c 2 -b -z -j -T -P 443 74.139.70.149" + args = ["mtr", "--report", "--show-ips", "--aslookup", "--json"] + if report_cycles is not None: + args.extend(["-c", str(int(report_cycles))]) + if port is not None: + if protocol is None: + protocol = IPProtocol.TCP + assert protocol in PROTOCOLS_W_PORT, "port only allowed for TCP/SCTP/UDP traces" + args.extend(["--port", str(int(port))]) + if protocol: + assert protocol in SUPPORTED_PROTOCOLS, f"unsupported protocol: {protocol}" + # default is ICMP (no args) + arg_map = { + IPProtocol.TCP: "--tcp", + IPProtocol.UDP: "--udp", + IPProtocol.SCTP: "--sctp", + } + if protocol in arg_map: + args.append(arg_map[protocol]) + args.append(ip) + return " ".join(args) + + +def get_mtr_version() -> str: + proc = subprocess.run( + ["mtr", "-v"], + capture_output=True, + text=True, + check=False, + ) + # e.g. mtr 0.95 + ver_str = proc.stdout.strip() + return ver_str.split(" ", 1)[1] + + +def run_mtr( + ip: str, + protocol: Optional[IPProtocol] = None, + port: Optional[int] = None, + report_cycles: int = 10, +) -> MTRResult: + args = build_mtr_command( + ip=ip, protocol=protocol, port=port, report_cycles=report_cycles + ) + proc = subprocess.run( + args.split(" "), + capture_output=True, + text=True, + check=False, + ) + raw = proc.stdout.strip() + return parse_mtr_output(raw, protocol=protocol, port=port) diff --git a/generalresearch/models/network/mtr/execute.py b/generalresearch/models/network/mtr/execute.py new file mode 100644 index 0000000..bd556bc --- /dev/null +++ b/generalresearch/models/network/mtr/execute.py @@ -0,0 +1,52 @@ +from datetime import datetime, timezone +from typing import Optional +from uuid import uuid4 + +from generalresearch.models.custom_types import UUIDStr +from generalresearch.models.network.definitions import IPProtocol +from generalresearch.models.network.mtr.command import ( + run_mtr, + get_mtr_version, + build_mtr_command, +) +from generalresearch.models.network.tool_run import MTRRun, ToolName, ToolClass, Status +from generalresearch.models.network.tool_run_command import ToolRunCommand +from generalresearch.models.network.utils import get_source_ip + + +def execute_mtr( + ip: str, + scan_group_id: Optional[UUIDStr] = None, + protocol: Optional[IPProtocol] = None, + port: Optional[int] = None, + report_cycles: int = 10, +) -> MTRRun: + started_at = datetime.now(tz=timezone.utc) + tool_version = get_mtr_version() + result = run_mtr(ip, protocol=protocol, port=port, report_cycles=report_cycles) + finished_at = datetime.now(tz=timezone.utc) + raw_command = build_mtr_command(ip) + config = ToolRunCommand( + command="mtr", + options={ + "protocol": protocol, + "port": port, + "report_cycles": report_cycles, + }, + ) + + return MTRRun( + tool_name=ToolName.MTR, + tool_class=ToolClass.TRACEROUTE, + tool_version=tool_version, + status=Status.SUCCESS, + ip=ip, + started_at=started_at, + finished_at=finished_at, + raw_command=raw_command, + scan_group_id=scan_group_id or uuid4().hex, + config=config, + parsed=result, + source_ip=get_source_ip(), + facility_id=1, + ) diff --git a/generalresearch/models/network/mtr/features.py b/generalresearch/models/network/mtr/features.py new file mode 100644 index 0000000..e7f2ff1 --- /dev/null +++ b/generalresearch/models/network/mtr/features.py @@ -0,0 +1,146 @@ +from typing import List, Dict + +from pydantic import BaseModel, ConfigDict, Field + +from generalresearch.models.network.definitions import IPKind +from generalresearch.models.network.mtr import MTRHop + + +class MTRFeatures(BaseModel): + model_config = ConfigDict() + + hop_count: int = Field() + + public_hop_count: int + private_hop_count: int + + unique_asn_count: int + asn_transition_count: int + + missing_hop_count: int + missing_hop_ratio: float + + # typical for mobile (vs residential) + private_hops_after_public: int + + has_cgnat: bool + + +def trim_local_hops(hops: List[MTRHop]) -> List[MTRHop]: + start = 0 + for i, h in enumerate(hops): + if h.ip_kind == IPKind.PUBLIC: + start = i + break + return hops[start:] + + +def extract_mtr_features(hops: List[MTRHop]) -> Dict[str, float | int | bool | None]: + features: Dict[str, float | int | bool | None] = {} + + if not hops: + return {} + + hops = trim_local_hops(hops) + + features["hop_count"] = len(hops) + + private_hops = 0 + public_hops = 0 + for h in hops: + if not h.ip: + continue + if h.ip_kind == IPKind.PUBLIC: + public_hops += 1 + else: + private_hops += 1 + features["private_hop_count"] = private_hops + features["public_hop_count"] = public_hops + + # ----------------------- + # ASN structure + # ----------------------- + + asns = [h.asn for h in hops if h.asn] + + features["unique_asn_count"] = len(set(asns)) + + asn_changes = 0 + for a, b in zip(asns, asns[1:]): + if a != b: + asn_changes += 1 + + features["asn_transition_count"] = asn_changes + + # ----------------------- + # Missing hops + # ----------------------- + + missing_hops = sum(1 for h in hops if h.ip is None) + + features["missing_hop_count"] = missing_hops + features["missing_hop_ratio"] = missing_hops / len(hops) + + # ----------------------- + # Packet loss + # ----------------------- + + lossy_hops = sum(1 for h in hops if h.loss_pct > 0) + + features["lossy_hop_count"] = lossy_hops + features["max_loss_pct"] = max(h.loss_pct for h in hops) + + # ----------------------- + # Latency stats + # ----------------------- + + avg_rtts = [h.avg_ms for h in hops if h.avg_ms > 0] + + if avg_rtts: + features["destination_rtt"] = avg_rtts[-1] + features["mean_rtt"] = sum(avg_rtts) / len(avg_rtts) + features["max_rtt"] = max(avg_rtts) + else: + features["destination_rtt"] = None + features["mean_rtt"] = None + features["max_rtt"] = None + + # ----------------------- + # RTT jumps + # ----------------------- + + rtt_jumps = [] + + for a, b in zip(hops, hops[1:]): + if a.avg_ms > 0 and b.avg_ms > 0: + rtt_jumps.append(b.avg_ms - a.avg_ms) + + if rtt_jumps: + features["max_rtt_jump"] = max(rtt_jumps) + features["mean_rtt_jump"] = sum(rtt_jumps) / len(rtt_jumps) + else: + features["max_rtt_jump"] = None + features["mean_rtt_jump"] = None + + # ----------------------- + # Jitter + # ----------------------- + + stdevs = [h.stdev_ms for h in hops if h.stdev_ms > 0] + + if stdevs: + features["max_jitter"] = max(stdevs) + features["mean_jitter"] = sum(stdevs) / len(stdevs) + else: + features["max_jitter"] = None + features["mean_jitter"] = None + + # ----------------------- + # Route completion + # ----------------------- + + last = hops[-1] + + features["destination_reached"] = last.ip is not None and last.loss_pct < 100 + + return features diff --git a/generalresearch/models/network/mtr/parser.py b/generalresearch/models/network/mtr/parser.py new file mode 100644 index 0000000..dc108d9 --- /dev/null +++ b/generalresearch/models/network/mtr/parser.py @@ -0,0 +1,18 @@ +import json +from typing import Dict + +from generalresearch.models.network.mtr.result import MTRResult + + +def parse_mtr_output(raw: str, port, protocol) -> MTRResult: + data = parse_mtr_raw_output(raw) + data["port"] = port + data["protocol"] = protocol + return MTRResult.model_validate(data) + + +def parse_mtr_raw_output(raw: str) -> Dict: + data = json.loads(raw)["report"] + data.update(data.pop("mtr")) + data["hops"] = data.pop("hubs") + return data diff --git a/generalresearch/models/network/mtr.py b/generalresearch/models/network/mtr/result.py index 4b040de..62f92ab 100644 --- a/generalresearch/models/network/mtr.py +++ b/generalresearch/models/network/mtr/result.py @@ -1,9 +1,7 @@ -import json import re -import subprocess from functools import cached_property from ipaddress import ip_address -from typing import List, Optional, Dict +from typing import List, Optional import tldextract from pydantic import ( @@ -17,6 +15,8 @@ from pydantic import ( from generalresearch.models.network.definitions import IPProtocol, get_ip_kind, IPKind +HOST_RE = re.compile(r"^(?P<hostname>.+?) \((?P<ip>[^)]+)\)$") + class MTRHop(BaseModel): model_config = ConfigDict(populate_by_name=True) @@ -105,7 +105,7 @@ class MTRHop(BaseModel): return d -class MTRReport(BaseModel): +class MTRResult(BaseModel): model_config = ConfigDict(populate_by_name=True) source: str = Field(description="Hostname of the system running mtr.", alias="src") @@ -165,85 +165,3 @@ class MTRReport(BaseModel): f"{hop.worst_ms:7.1f} " f"{hop.stdev_ms:7.1f}" ) - - -HOST_RE = re.compile(r"^(?P<hostname>.+?) \((?P<ip>[^)]+)\)$") - -SUPPORTED_PROTOCOLS = { - IPProtocol.TCP, - IPProtocol.UDP, - IPProtocol.SCTP, - IPProtocol.ICMP, -} -PROTOCOLS_W_PORT = {IPProtocol.TCP, IPProtocol.UDP, IPProtocol.SCTP} - - -def get_mtr_command( - ip: str, - protocol: Optional[IPProtocol] = None, - port: Optional[int] = None, - report_cycles: int = 10, -) -> List[str]: - # https://manpages.ubuntu.com/manpages/focal/man8/mtr.8.html - # e.g. "mtr -r -c 2 -b -z -j -T -P 443 74.139.70.149" - args = ["mtr", "--report", "--show-ips", "--aslookup", "--json"] - if report_cycles is not None: - args.extend(["-c", str(int(report_cycles))]) - if port is not None: - if protocol is None: - protocol = IPProtocol.TCP - assert protocol in PROTOCOLS_W_PORT, "port only allowed for TCP/SCTP/UDP traces" - args.extend(["--port", str(int(port))]) - if protocol: - assert protocol in SUPPORTED_PROTOCOLS, f"unsupported protocol: {protocol}" - # default is ICMP (no args) - arg_map = { - IPProtocol.TCP: "--tcp", - IPProtocol.UDP: "--udp", - IPProtocol.SCTP: "--sctp", - } - if protocol in arg_map: - args.append(arg_map[protocol]) - args.append(ip) - return args - - -def get_mtr_version() -> str: - proc = subprocess.run( - ["mtr", "-v"], - capture_output=True, - text=True, - check=False, - ) - # e.g. mtr 0.95 - ver_str = proc.stdout.strip() - return ver_str.split(" ", 1)[1] - - -def run_mtr( - ip: str, - protocol: Optional[IPProtocol] = None, - port: Optional[int] = None, - report_cycles: int = 10, -) -> MTRReport: - args = get_mtr_command( - ip=ip, protocol=protocol, port=port, report_cycles=report_cycles - ) - proc = subprocess.run( - args, - capture_output=True, - text=True, - check=False, - ) - raw = proc.stdout.strip() - data = parse_raw_output(raw) - data["port"] = port - data["protocol"] = protocol - return MTRReport.model_validate(data) - - -def parse_raw_output(raw: str) -> Dict: - data = json.loads(raw)["report"] - data.update(data.pop("mtr")) - data["hops"] = data.pop("hubs") - return data diff --git a/generalresearch/models/network/nmap/__init__.py b/generalresearch/models/network/nmap/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/generalresearch/models/network/nmap/__init__.py diff --git a/generalresearch/models/network/nmap/command.py b/generalresearch/models/network/nmap/command.py new file mode 100644 index 0000000..dfa55de --- /dev/null +++ b/generalresearch/models/network/nmap/command.py @@ -0,0 +1,25 @@ +import subprocess +from typing import Optional, List + +from generalresearch.models.network.nmap.parser import parse_nmap_xml +from generalresearch.models.network.nmap.result import NmapResult + + +def build_nmap_command(ip: str, top_ports: Optional[int] = 1000) -> List[str]: + # e.g. "nmap -Pn -T4 -A --top-ports 1000 -oX - scanme.nmap.org" + # https://linux.die.net/man/1/nmap + args = ["nmap", "-Pn", "-T4", "-A", "--top-ports", str(int(top_ports)), "-oX", "-"] + args.append(ip) + return args + + +def run_nmap(ip: str, top_ports: Optional[int] = 1000) -> NmapResult: + args = build_nmap_command(ip=ip, top_ports=top_ports) + proc = subprocess.run( + args, + capture_output=True, + text=True, + check=False, + ) + raw = proc.stdout.strip() + return parse_nmap_xml(raw) diff --git a/generalresearch/models/network/nmap/execute.py b/generalresearch/models/network/nmap/execute.py new file mode 100644 index 0000000..fc1e2fa --- /dev/null +++ b/generalresearch/models/network/nmap/execute.py @@ -0,0 +1,30 @@ +from typing import Optional +from uuid import uuid4 + +from generalresearch.models.custom_types import UUIDStr +from generalresearch.models.network.nmap.command import run_nmap +from generalresearch.models.network.tool_run import NmapRun, ToolName, ToolClass, Status +from generalresearch.models.network.tool_run_command import ToolRunCommand + + +def execute_nmap( + ip: str, top_ports: Optional[int] = 1000, scan_group_id: Optional[UUIDStr] = None +): + result = run_nmap(ip=ip, top_ports=top_ports) + assert result.exit_status == "success" + assert result.target_ip == ip + + run = NmapRun( + tool_name=ToolName.NMAP, + tool_class=ToolClass.PORT_SCAN, + tool_version=result.version, + status=Status.SUCCESS, + ip=ip, + started_at=result.started_at, + finished_at=result.finished_at, + raw_command=result.command_line, + scan_group_id=scan_group_id or uuid4().hex, + config=ToolRunCommand(command="nmap", options={'top_ports': top_ports}), + parsed=result, + ) + return run
\ No newline at end of file diff --git a/generalresearch/models/network/xml_parser.py b/generalresearch/models/network/nmap/parser.py index 349bc94..5a441bb 100644 --- a/generalresearch/models/network/xml_parser.py +++ b/generalresearch/models/network/nmap/parser.py @@ -3,9 +3,9 @@ from datetime import datetime, timezone from typing import List, Dict, Any, Tuple, Optional from generalresearch.models.network.definitions import IPProtocol -from generalresearch.models.network.nmap import ( +from generalresearch.models.network.nmap.result import ( NmapHostname, - NmapRun, + NmapResult, NmapPort, PortState, PortStateReason, @@ -39,7 +39,7 @@ class NmapXmlParser: """ @classmethod - def parse_xml(cls, nmap_data: str) -> NmapRun: + def parse_xml(cls, nmap_data: str) -> NmapResult: """ Expects a full nmap scan report. """ @@ -55,7 +55,7 @@ class NmapXmlParser: return cls._parse_xml_nmaprun(root) @classmethod - def _parse_xml_nmaprun(cls, root: ET.Element) -> NmapRun: + def _parse_xml_nmaprun(cls, root: ET.Element) -> NmapResult: """ This method parses out a full nmap scan report from its XML root node: <nmaprun>. We expect there is only 1 host in this report! @@ -79,7 +79,7 @@ class NmapXmlParser: nmap_data.update(cls._parse_xml_host(root.find(".//host"))) - return NmapRun.model_validate(nmap_data) + return NmapResult.model_validate(nmap_data) @classmethod def _validate_nmap_root(cls, root: ET.Element) -> None: @@ -406,3 +406,7 @@ class NmapXmlParser: protocol=IPProtocol(proto_attr) if proto_attr is not None else None, hops=hops, ) + + +def parse_nmap_xml(raw): + return NmapXmlParser.parse_xml(raw) diff --git a/generalresearch/models/network/nmap.py b/generalresearch/models/network/nmap/result.py index 4656a93..635db06 100644 --- a/generalresearch/models/network/nmap.py +++ b/generalresearch/models/network/nmap/result.py @@ -1,5 +1,4 @@ import json -import subprocess from datetime import timedelta from enum import StrEnum from functools import cached_property @@ -271,7 +270,7 @@ class NmapScanInfo(BaseModel): return ports -class NmapRun(BaseModel): +class NmapResult(BaseModel): """ A Nmap Run. Expects that we've only scanned ONE host. """ @@ -430,28 +429,3 @@ class NmapRun(BaseModel): d["parsed"] = self.model_dump_json(indent=0) d["open_tcp_ports"] = json.dumps(self.tcp_open_ports) return d - - -def get_nmap_command(ip: str, top_ports: Optional[int] = 1000) -> List[str]: - # e.g. "nmap -Pn -T4 -A --top-ports 1000 -oX - scanme.nmap.org" - # https://linux.die.net/man/1/nmap - args = ["nmap", "-Pn", "-T4", "-A", "--top-ports", str(int(top_ports)), "-oX", "-"] - args.append(ip) - return args - - -def run_nmap(ip: str, top_ports: Optional[int] = 1000) -> NmapRun: - from generalresearch.models.network.xml_parser import NmapXmlParser - - p = NmapXmlParser() - - args = get_nmap_command(ip=ip, top_ports=top_ports) - proc = subprocess.run( - args, - capture_output=True, - text=True, - check=False, - ) - raw = proc.stdout.strip() - n = p.parse_xml(raw) - return n diff --git a/generalresearch/models/network/rdns/__init__.py b/generalresearch/models/network/rdns/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/generalresearch/models/network/rdns/__init__.py diff --git a/generalresearch/models/network/rdns/command.py b/generalresearch/models/network/rdns/command.py new file mode 100644 index 0000000..e9d5bfd --- /dev/null +++ b/generalresearch/models/network/rdns/command.py @@ -0,0 +1,33 @@ +import subprocess + +from generalresearch.models.network.rdns.parser import parse_rdns_output +from generalresearch.models.network.rdns.result import RDNSResult + + +def run_rdns(ip: str) -> RDNSResult: + args = build_rdns_command(ip).split(" ") + proc = subprocess.run( + args, + capture_output=True, + text=True, + check=False, + ) + raw = proc.stdout.strip() + return parse_rdns_output(ip, raw) + + +def build_rdns_command(ip: str): + # e.g. dig +noall +answer -x 1.2.3.4 + return " ".join(["dig", "+noall", "+answer", "-x", ip]) + + +def get_dig_version() -> str: + proc = subprocess.run( + ["dig", "-v"], + capture_output=True, + text=True, + check=False, + ) + # e.g. DiG 9.18.39-0ubuntu0.22.04.2-Ubuntu + ver_str = proc.stderr.strip() + return ver_str.split("-", 1)[0].split(" ", 1)[1] diff --git a/generalresearch/models/network/rdns/execute.py b/generalresearch/models/network/rdns/execute.py new file mode 100644 index 0000000..97b8bf8 --- /dev/null +++ b/generalresearch/models/network/rdns/execute.py @@ -0,0 +1,41 @@ +from datetime import datetime, timezone +from typing import Optional +from uuid import uuid4 + +from generalresearch.models.custom_types import UUIDStr +from generalresearch.models.network.rdns.command import ( + run_rdns, + get_dig_version, + build_rdns_command, +) +from generalresearch.models.network.tool_run import ( + ToolName, + ToolClass, + Status, + RDNSRun, +) +from generalresearch.models.network.tool_run_command import ToolRunCommand + + +def execute_rdns(ip: str, scan_group_id: Optional[UUIDStr] = None): + started_at = datetime.now(tz=timezone.utc) + tool_version = get_dig_version() + result = run_rdns(ip=ip) + finished_at = datetime.now(tz=timezone.utc) + raw_command = build_rdns_command(ip) + + run = RDNSRun( + tool_name=ToolName.DIG, + tool_class=ToolClass.RDNS, + tool_version=tool_version, + status=Status.SUCCESS, + ip=ip, + started_at=started_at, + finished_at=finished_at, + raw_command=raw_command, + scan_group_id=scan_group_id or uuid4().hex, + config=ToolRunCommand(command="dig", options={}), + parsed=result, + ) + + return run diff --git a/generalresearch/models/network/rdns/parser.py b/generalresearch/models/network/rdns/parser.py new file mode 100644 index 0000000..f12a6f4 --- /dev/null +++ b/generalresearch/models/network/rdns/parser.py @@ -0,0 +1,21 @@ +import ipaddress +import re +from typing import List + +from generalresearch.models.network.rdns.result import RDNSResult + +PTR_RE = re.compile(r"\sPTR\s+([^\s]+)\.") + + +def parse_rdns_output(ip, raw): + hostnames: List[str] = [] + + for line in raw.splitlines(): + m = PTR_RE.search(line) + if m: + hostnames.append(m.group(1)) + + return RDNSResult( + ip=ipaddress.ip_address(ip), + hostnames=hostnames, + ) diff --git a/generalresearch/models/network/rdns.py b/generalresearch/models/network/rdns/result.py index e00a32d..81b4085 100644 --- a/generalresearch/models/network/rdns.py +++ b/generalresearch/models/network/rdns/result.py @@ -1,18 +1,11 @@ -import ipaddress import json from functools import cached_property - -from pydantic import BaseModel, Field, model_validator, computed_field from typing import Optional, List -from typing_extensions import Self +import tldextract +from pydantic import BaseModel, Field, model_validator, computed_field from generalresearch.models.custom_types import IPvAnyAddressStr -import subprocess -import re -from typing import List -import ipaddress -import tldextract class RDNSResult(BaseModel): @@ -44,7 +37,9 @@ class RDNSResult(BaseModel): @cached_property def primary_domain(self) -> Optional[str]: if self.primary_hostname: - return tldextract.extract(self.primary_hostname).top_domain_under_public_suffix + return tldextract.extract( + self.primary_hostname + ).top_domain_under_public_suffix def model_dump_postgres(self): # Writes for the network_rdnsresult table @@ -54,48 +49,3 @@ class RDNSResult(BaseModel): ) d["hostnames"] = json.dumps(self.hostnames) return d - - @classmethod - def from_dig(cls, ip: str, raw_output: str) -> Self: - hostnames: List[str] = [] - - for line in raw_output.splitlines(): - m = PTR_RE.search(line) - if m: - hostnames.append(m.group(1)) - - return cls( - ip=ipaddress.ip_address(ip), - hostnames=hostnames, - ) - - -PTR_RE = re.compile(r"\sPTR\s+([^\s]+)\.") - - -def dig_rdns(ip: str) -> RDNSResult: - args = get_dig_rdns_command(ip).split(" ") - proc = subprocess.run( - args, - capture_output=True, - text=True, - check=False, - ) - raw = proc.stdout.strip() - return RDNSResult.from_dig(ip=ip, raw_output=raw) - - -def get_dig_rdns_command(ip: str): - return " ".join(["dig", "+noall", "+answer", "-x", ip]) - - -def get_dig_version() -> str: - proc = subprocess.run( - ["dig", "-v"], - capture_output=True, - text=True, - check=False, - ) - # e.g. DiG 9.18.39-0ubuntu0.22.04.2-Ubuntu - ver_str = proc.stderr.strip() - return ver_str.split("-", 1)[0].split(" ", 1)[1] diff --git a/generalresearch/models/network/tool_run.py b/generalresearch/models/network/tool_run.py index 2588890..36e6950 100644 --- a/generalresearch/models/network/tool_run.py +++ b/generalresearch/models/network/tool_run.py @@ -1,6 +1,5 @@ -from datetime import datetime, timezone from enum import StrEnum -from typing import Optional, Tuple +from typing import Optional, Literal from uuid import uuid4 from pydantic import BaseModel, Field, PositiveInt @@ -10,20 +9,10 @@ from generalresearch.models.custom_types import ( IPvAnyAddressStr, UUIDStr, ) -from generalresearch.models.network.nmap import NmapRun -from generalresearch.models.network.rdns import ( - RDNSResult, - get_dig_version, - dig_rdns, - get_dig_rdns_command, -) -from generalresearch.models.network.mtr import ( - MTRReport, - get_mtr_version, - run_mtr, - get_mtr_command, -) -from generalresearch.models.network.tool_utils import ToolRunCommand +from generalresearch.models.network.mtr.result import MTRResult +from generalresearch.models.network.nmap.result import NmapResult +from generalresearch.models.network.rdns.result import RDNSResult +from generalresearch.models.network.tool_run_command import ToolRunCommand class ToolClass(StrEnum): @@ -76,8 +65,11 @@ class ToolRun(BaseModel): return d -class PortScanRun(ToolRun): - parsed: NmapRun = Field() +class NmapRun(ToolRun): + tool_class: Literal[ToolClass.PORT_SCAN] = Field(default=ToolClass.PORT_SCAN) + tool_name: Literal[ToolName.NMAP] = Field(default=ToolName.NMAP) + + parsed: NmapResult = Field() def model_dump_postgres(self): d = super().model_dump_postgres() @@ -86,7 +78,10 @@ class PortScanRun(ToolRun): return d -class RDnsRun(ToolRun): +class RDNSRun(ToolRun): + tool_class: Literal[ToolClass.RDNS] = Field(default=ToolClass.RDNS) + tool_name: Literal[ToolName.DIG] = Field(default=ToolName.DIG) + parsed: RDNSResult = Field() def model_dump_postgres(self): @@ -96,10 +91,13 @@ class RDnsRun(ToolRun): return d -class MtrRun(ToolRun): +class MTRRun(ToolRun): + tool_class: Literal[ToolClass.TRACEROUTE] = Field(default=ToolClass.TRACEROUTE) + tool_name: Literal[ToolName.MTR] = Field(default=ToolName.MTR) + facility_id: int = Field(default=1) source_ip: IPvAnyAddressStr = Field() - parsed: MTRReport = Field() + parsed: MTRResult = Field() def model_dump_postgres(self): d = super().model_dump_postgres() @@ -108,66 +106,3 @@ class MtrRun(ToolRun): d["facility_id"] = self.facility_id d.update(self.parsed.model_dump_postgres()) return d - - -def new_tool_run_from_nmap( - nmap_run: NmapRun, scan_group_id: Optional[UUIDStr] = None -) -> PortScanRun: - assert nmap_run.exit_status == "success" - return PortScanRun( - tool_name=ToolName.NMAP, - tool_class=ToolClass.PORT_SCAN, - tool_version=nmap_run.version, - status=Status.SUCCESS, - ip=nmap_run.target_ip, - started_at=nmap_run.started_at, - finished_at=nmap_run.finished_at, - raw_command=nmap_run.command_line, - scan_group_id=scan_group_id or uuid4().hex, - config=ToolRunCommand.from_raw_command(nmap_run.command_line), - parsed=nmap_run, - ) - - -def run_dig(ip: str, scan_group_id: Optional[UUIDStr] = None) -> RDnsRun: - started_at = datetime.now(tz=timezone.utc) - tool_version = get_dig_version() - rdns_result = dig_rdns(ip) - finished_at = datetime.now(tz=timezone.utc) - raw_command = get_dig_rdns_command(ip) - - return RDnsRun( - tool_name=ToolName.DIG, - tool_class=ToolClass.RDNS, - tool_version=tool_version, - status=Status.SUCCESS, - ip=ip, - started_at=started_at, - finished_at=finished_at, - raw_command=raw_command, - scan_group_id=scan_group_id or uuid4().hex, - config=ToolRunCommand.from_raw_command(raw_command), - parsed=rdns_result, - ) - - -def mtr_tool_run(ip: str, scan_group_id: Optional[UUIDStr] = None) -> MtrRun: - started_at = datetime.now(tz=timezone.utc) - tool_version = get_mtr_version() - result = run_mtr(ip) - finished_at = datetime.now(tz=timezone.utc) - raw_command = " ".join(get_mtr_command(ip)) - - return MtrRun( - tool_name=ToolName.MTR, - tool_class=ToolClass.TRACEROUTE, - tool_version=tool_version, - status=Status.SUCCESS, - ip=ip, - started_at=started_at, - finished_at=finished_at, - raw_command=raw_command, - scan_group_id=scan_group_id or uuid4().hex, - config=ToolRunCommand.from_raw_command(raw_command), - parsed=result, - ) diff --git a/generalresearch/models/network/tool_run_command.py b/generalresearch/models/network/tool_run_command.py new file mode 100644 index 0000000..e3d94df --- /dev/null +++ b/generalresearch/models/network/tool_run_command.py @@ -0,0 +1,9 @@ +from typing import Dict + +from pydantic import BaseModel, Field + + +class ToolRunCommand(BaseModel): + # todo: expand with arguments specific for each tool + command: str = Field() + options: Dict[str, str | int] = Field(default_factory=dict) diff --git a/generalresearch/models/network/tool_utils.py b/generalresearch/models/network/tool_utils.py deleted file mode 100644 index 83d988d..0000000 --- a/generalresearch/models/network/tool_utils.py +++ /dev/null @@ -1,69 +0,0 @@ -import shlex -from typing import Dict, List - -from pydantic import BaseModel -from typing_extensions import Self - -""" -e.g.: "nmap -Pn -sV -p 80,443 --reason --max-retries=3 1.2.3.4" -{'command': 'nmap', - 'options': {'p': '80,443', 'max-retries': '3'}, - 'flags': ['Pn', 'sV', 'reason'], - 'positionals': ['1.2.3.4']} -""" - - -class ToolRunCommand(BaseModel): - command: str - options: Dict[str, str] - flags: List[str] - positionals: List[str] - - @classmethod - def from_raw_command(cls, s: str) -> Self: - return cls.model_validate(parse_command(s)) - - -def parse_command(cmd: str): - tokens = shlex.split(cmd) - - result = { - "command": tokens[0], - "options": {}, - "flags": [], - "positionals": [], - } - - i = 1 - while i < len(tokens): - tok = tokens[i] - - # --key=value - if tok.startswith("--") and "=" in tok: - k, v = tok[2:].split("=", 1) - result["options"][k] = v - - # --key value - elif tok.startswith("--"): - key = tok[2:] - if i + 1 < len(tokens) and not tokens[i + 1].startswith("-"): - result["options"][key] = tokens[i + 1] - i += 1 - else: - result["flags"].append(key) - - # short flag or short flag with arg - elif tok.startswith("-"): - if i + 1 < len(tokens) and not tokens[i + 1].startswith("-"): - result["options"][tok[1:]] = tokens[i + 1] - i += 1 - else: - result["flags"].append(tok[1:]) - - else: - result["positionals"].append(tok) - - i += 1 - - result["flags"] = sorted(result["flags"]) - return result diff --git a/generalresearch/models/network/utils.py b/generalresearch/models/network/utils.py new file mode 100644 index 0000000..fee9b80 --- /dev/null +++ b/generalresearch/models/network/utils.py @@ -0,0 +1,5 @@ +import requests + + +def get_source_ip(): + return requests.get("https://icanhazip.com?").text.strip() diff --git a/test_utils/conftest.py b/test_utils/conftest.py index 187ff58..0e712bb 100644 --- a/test_utils/conftest.py +++ b/test_utils/conftest.py @@ -4,6 +4,7 @@ from os.path import join as pjoin from pathlib import Path from typing import TYPE_CHECKING, Callable from uuid import uuid4 +from datetime import datetime, timedelta, timezone import pytest import redis @@ -17,8 +18,6 @@ from generalresearch.redis_helper import RedisConfig from generalresearch.sql_helper import SqlHelper if TYPE_CHECKING: - from datetime import datetime - from generalresearch.config import GRLBaseSettings from generalresearch.currency import USDCent from generalresearch.models.thl.session import Status @@ -38,6 +37,7 @@ def env_file_path(pytestconfig: Config) -> str: @pytest.fixture(scope="session") def settings(env_file_path: str) -> "GRLBaseSettings": from generalresearch.config import GRLBaseSettings + print(f"{env_file_path=}") s = GRLBaseSettings(_env_file=env_file_path) @@ -203,16 +203,12 @@ def wall_status(request) -> "Status": @pytest.fixture -def utc_now() -> "datetime": - from datetime import datetime, timezone - +def utc_now() -> datetime: return datetime.now(tz=timezone.utc) @pytest.fixture -def utc_hour_ago() -> "datetime": - from datetime import datetime, timedelta, timezone - +def utc_hour_ago() -> datetime: return datetime.now(tz=timezone.utc) - timedelta(hours=1) diff --git a/test_utils/managers/network/conftest.py b/test_utils/managers/network/conftest.py index 70fda4e..f6a4078 100644 --- a/test_utils/managers/network/conftest.py +++ b/test_utils/managers/network/conftest.py @@ -1,40 +1,31 @@ import os -from datetime import datetime, timezone -from typing import Callable, TYPE_CHECKING +from datetime import timedelta, datetime, timezone from uuid import uuid4 import pytest from generalresearch.managers.network.label import IPLabelManager -from generalresearch.managers.network.nmap import NmapManager from generalresearch.managers.network.tool_run import ToolRunManager -from generalresearch.models.network.rdns import ( - RDNSResult, - get_dig_version, - get_dig_rdns_command, -) -from generalresearch.models.network.tool_run import ( - RDnsRun, - ToolName, - ToolClass, - Status, -) -from generalresearch.models.network.tool_utils import ToolRunCommand -from generalresearch.models.network.xml_parser import NmapXmlParser +from generalresearch.models.network.definitions import IPProtocol +from generalresearch.models.network.mtr.command import build_mtr_command +from generalresearch.models.network.mtr.parser import parse_mtr_output +from generalresearch.models.network.nmap.parser import parse_nmap_xml +from generalresearch.models.network.rdns.command import build_rdns_command +from generalresearch.models.network.rdns.parser import parse_rdns_output +from generalresearch.models.network.tool_run import NmapRun, Status, RDNSRun, MTRRun +from generalresearch.models.network.tool_run_command import ToolRunCommand @pytest.fixture(scope="session") -def iplabel_manager(thl_web_rw) -> IPLabelManager: - assert "/unittest-" in thl_web_rw.dsn.path - - return IPLabelManager(pg_config=thl_web_rw) +def scan_group_id(): + return uuid4().hex @pytest.fixture(scope="session") -def nmap_manager(thl_web_rw) -> NmapManager: +def iplabel_manager(thl_web_rw) -> IPLabelManager: assert "/unittest-" in thl_web_rw.dsn.path - return NmapManager(pg_config=thl_web_rw) + return IPLabelManager(pg_config=thl_web_rw) @pytest.fixture(scope="session") @@ -45,7 +36,7 @@ def toolrun_manager(thl_web_rw) -> ToolRunManager: @pytest.fixture(scope="session") -def nmap_xml_str(request) -> str: +def nmap_raw_output(request) -> str: fp = os.path.join(request.config.rootpath, "data/nmaprun1.xml") with open(fp, "r") as f: data = f.read() @@ -53,34 +44,84 @@ def nmap_xml_str(request) -> str: @pytest.fixture(scope="session") -def nmap_run(nmap_xml_str): - return NmapXmlParser.parse_xml(nmap_xml_str) +def nmap_result(nmap_raw_output): + return parse_nmap_xml(nmap_raw_output) @pytest.fixture(scope="session") -def raw_dig_output(): +def nmap_run(nmap_result, scan_group_id): + r = nmap_result + return NmapRun( + tool_version=r.version, + status=Status.SUCCESS, + ip=r.target_ip, + started_at=r.started_at, + finished_at=r.finished_at, + raw_command=r.command_line, + scan_group_id=scan_group_id, + config=ToolRunCommand(command="nmap"), + parsed=r, + ) + + +@pytest.fixture(scope="session") +def dig_raw_output(): return "156.32.33.45.in-addr.arpa. 300 IN PTR scanme.nmap.org." @pytest.fixture(scope="session") -def reverse_dns_run(raw_dig_output): +def rdns_result(dig_raw_output): + return parse_rdns_output(ip="45.33.32.156", raw=dig_raw_output) + + +@pytest.fixture(scope="session") +def rdns_run(rdns_result, scan_group_id): + r = rdns_result ip = "45.33.32.156" - rdns_result = RDNSResult.from_dig(ip=ip, raw_output=raw_dig_output) - scan_group_id = uuid4().hex - started_at = datetime.now(tz=timezone.utc) - tool_version = get_dig_version() - finished_at = datetime.now(tz=timezone.utc) - raw_command = get_dig_rdns_command(ip) - return RDnsRun( - tool_name=ToolName.DIG, - tool_class=ToolClass.RDNS, - tool_version=tool_version, + utc_now = datetime.now(tz=timezone.utc) + return RDNSRun( + tool_version="1.2.3", status=Status.SUCCESS, ip=ip, - started_at=started_at, - finished_at=finished_at, - raw_command=raw_command, - scan_group_id=scan_group_id or uuid4().hex, - config=ToolRunCommand.from_raw_command(raw_command), - parsed=rdns_result, + started_at=utc_now, + finished_at=utc_now + timedelta(seconds=1), + raw_command=build_rdns_command(ip=ip), + scan_group_id=scan_group_id, + config=ToolRunCommand(command="dig"), + parsed=r, + ) + + +@pytest.fixture(scope="session") +def mtr_raw_output(request): + fp = os.path.join(request.config.rootpath, "data/mtr_fatbeam.json") + with open(fp, "r") as f: + data = f.read() + return data + + +@pytest.fixture(scope="session") +def mtr_result(mtr_raw_output): + return parse_mtr_output(mtr_raw_output, port=443, protocol=IPProtocol.TCP) + + +@pytest.fixture(scope="session") +def mtr_run(mtr_result, scan_group_id): + r = mtr_result + utc_now = datetime.now(tz=timezone.utc) + + return MTRRun( + tool_version="1.2.3", + status=Status.SUCCESS, + ip=r.destination, + started_at=utc_now, + finished_at=utc_now + timedelta(seconds=1), + raw_command=build_mtr_command( + ip=r.destination, protocol=IPProtocol.TCP, port=443 + ), + scan_group_id=scan_group_id, + config=ToolRunCommand(command="mtr"), + parsed=r, + facility_id=1, + source_ip="1.2.3.4", ) diff --git a/tests/managers/network/tool_run.py b/tests/managers/network/tool_run.py index 0f9388f..c05af92 100644 --- a/tests/managers/network/tool_run.py +++ b/tests/managers/network/tool_run.py @@ -6,102 +6,38 @@ import faker import pytest from generalresearch.models.network.definitions import IPProtocol -from generalresearch.models.network.mtr import ( - get_mtr_version, - parse_raw_output, - MTRReport, - get_mtr_command, -) + from generalresearch.models.network.tool_run import ( - new_tool_run_from_nmap, - run_dig, - MtrRun, ToolName, ToolClass, Status, ) -from generalresearch.models.network.tool_utils import ToolRunCommand fake = faker.Faker() -def test_create_tool_run_from_nmap(nmap_run, toolrun_manager): - scan_group_id = uuid4().hex - run = new_tool_run_from_nmap(nmap_run, scan_group_id=scan_group_id) - - toolrun_manager.create_portscan_run(run) - - run_out = toolrun_manager.get_portscan_run(run.id) - - assert run == run_out - - -def test_create_tool_run_from_dig_fixture(reverse_dns_run, toolrun_manager): - - toolrun_manager.create_rdns_run(reverse_dns_run) - - run_out = toolrun_manager.get_rdns_run(reverse_dns_run.id) - - assert reverse_dns_run == run_out - - -def test_run_dig(toolrun_manager): - reverse_dns_run = run_dig(ip="65.19.129.53") - - toolrun_manager.create_rdns_run(reverse_dns_run) - - run_out = toolrun_manager.get_rdns_run(reverse_dns_run.id) - - assert reverse_dns_run == run_out - - -def test_run_dig_empty(toolrun_manager): - reverse_dns_run = run_dig(ip=fake.ipv6()) +def test_create_tool_run_from_nmap_run(nmap_run, toolrun_manager): - toolrun_manager.create_rdns_run(reverse_dns_run) + toolrun_manager.create_nmap_run(nmap_run) - run_out = toolrun_manager.get_rdns_run(reverse_dns_run.id) + run_out = toolrun_manager.get_nmap_run(nmap_run.id) - assert reverse_dns_run == run_out + assert nmap_run == run_out -@pytest.fixture(scope="session") -def mtr_report(request) -> MTRReport: - fp = os.path.join(request.config.rootpath, "data/mtr_fatbeam.json") - with open(fp, "r") as f: - s = f.read() - data = parse_raw_output(s) - data["port"] = 443 - data["protocol"] = IPProtocol.TCP - return MTRReport.model_validate(data) +def test_create_tool_run_from_rdns_run(rdns_run, toolrun_manager): + toolrun_manager.create_rdns_run(rdns_run) -def test_create_tool_run_from_mtr(toolrun_manager, mtr_report): - started_at = datetime.now(tz=timezone.utc) - tool_version = get_mtr_version() + run_out = toolrun_manager.get_rdns_run(rdns_run.id) - ip = mtr_report.destination + assert rdns_run == run_out - finished_at = datetime.now(tz=timezone.utc) - raw_command = " ".join(get_mtr_command(ip)) - run = MtrRun( - tool_name=ToolName.MTR, - tool_class=ToolClass.TRACEROUTE, - tool_version=tool_version, - status=Status.SUCCESS, - ip=ip, - started_at=started_at, - finished_at=finished_at, - raw_command=raw_command, - scan_group_id=uuid4().hex, - config=ToolRunCommand.from_raw_command(raw_command), - parsed=mtr_report, - source_ip="1.1.1.1" - ) +def test_create_tool_run_from_mtr_run(mtr_run, toolrun_manager): - toolrun_manager.create_mtr_run(run) + toolrun_manager.create_mtr_run(mtr_run) - run_out = toolrun_manager.get_mtr_run(run.id) + run_out = toolrun_manager.get_mtr_run(mtr_run.id) - assert run == run_out + assert mtr_run == run_out diff --git a/tests/models/network/nmap.py b/tests/models/network/nmap.py deleted file mode 100644 index 4fc7014..0000000 --- a/tests/models/network/nmap.py +++ /dev/null @@ -1,32 +0,0 @@ -import os - -import pytest - -from generalresearch.models.network.xml_parser import NmapXmlParser - - -@pytest.fixture -def nmap_xml_str(request) -> str: - fp = os.path.join(request.config.rootpath, "data/nmaprun1.xml") - with open(fp, "r") as f: - data = f.read() - return data - - -@pytest.fixture -def nmap_xml_str2(request) -> str: - fp = os.path.join(request.config.rootpath, "data/nmaprun2.xml") - with open(fp, "r") as f: - data = f.read() - return data - - -def test_nmap_xml_parser(nmap_xml_str, nmap_xml_str2): - p = NmapXmlParser() - n = p.parse_xml(nmap_xml_str) - assert n.tcp_open_ports == [61232] - assert len(n.trace.hops) == 18 - - n = p.parse_xml(nmap_xml_str2) - assert n.tcp_open_ports == [22, 80, 9929, 31337] - assert n.trace is None diff --git a/tests/models/network/nmap_parser.py b/tests/models/network/nmap_parser.py new file mode 100644 index 0000000..96d7b37 --- /dev/null +++ b/tests/models/network/nmap_parser.py @@ -0,0 +1,22 @@ +import os + +import pytest + +from generalresearch.models.network.nmap.parser import parse_nmap_xml + +@pytest.fixture +def nmap_raw_output_2(request) -> str: + fp = os.path.join(request.config.rootpath, "data/nmaprun2.xml") + with open(fp, "r") as f: + data = f.read() + return data + + +def test_nmap_xml_parser(nmap_raw_output, nmap_raw_output_2): + n = parse_nmap_xml(nmap_raw_output) + assert n.tcp_open_ports == [61232] + assert len(n.trace.hops) == 18 + + n = parse_nmap_xml(nmap_raw_output_2) + assert n.tcp_open_ports == [22, 80, 9929, 31337] + assert n.trace is None diff --git a/tests/models/network/rdns.py b/tests/models/network/rdns.py index 9167749..64e8351 100644 --- a/tests/models/network/rdns.py +++ b/tests/models/network/rdns.py @@ -1,23 +1,45 @@ -from generalresearch.models.network.rdns import dig_rdns -import faker +# from generalresearch.models.network.rdns import run_rdns +# import faker +# +# fake = faker.Faker() +# +# +# def test_dig_rdns(): +# # Actually runs dig -x. Idk how stable this is +# ip = "45.33.32.156" +# rdns_result = run_rdns(ip) +# assert rdns_result.primary_hostname == "scanme.nmap.org" +# assert rdns_result.primary_org == "nmap" +# +# ip = "65.19.129.53" +# rdns_result = run_rdns(ip) +# assert rdns_result.primary_hostname == "in1-smtp.grlengine.com" +# assert rdns_result.primary_org == "grlengine" +# +# ip = fake.ipv6() +# rdns_result = run_rdns(ip) +# assert rdns_result.primary_hostname is None +# assert rdns_result.primary_org is None +# print(rdns_result.model_dump_postgres()) -fake = faker.Faker() - -def test_dig_rdns(): - # Actually runs dig -x. Idk how stable this is - ip = "45.33.32.156" - rdns_result = dig_rdns(ip) - assert rdns_result.primary_hostname == "scanme.nmap.org" - assert rdns_result.primary_org == "nmap" - - ip = "65.19.129.53" - rdns_result = dig_rdns(ip) - assert rdns_result.primary_hostname == "in1-smtp.grlengine.com" - assert rdns_result.primary_org == "grlengine" - - ip = fake.ipv6() - rdns_result = dig_rdns(ip) - assert rdns_result.primary_hostname is None - assert rdns_result.primary_org is None - print(rdns_result.model_dump_postgres()) +# +# +# def test_run_dig(toolrun_manager): +# reverse_dns_run = run_dig(ip="65.19.129.53") +# +# toolrun_manager.create_rdns_run(reverse_dns_run) +# +# run_out = toolrun_manager.get_rdns_run(reverse_dns_run.id) +# +# assert reverse_dns_run == run_out +# +# +# def test_run_dig_empty(toolrun_manager): +# reverse_dns_run = run_dig(ip=fake.ipv6()) +# +# toolrun_manager.create_rdns_run(reverse_dns_run) +# +# run_out = toolrun_manager.get_rdns_run(reverse_dns_run.id) +# +# assert reverse_dns_run == run_out
\ No newline at end of file diff --git a/tests/models/network/tool_run.py b/tests/models/network/tool_run.py deleted file mode 100644 index c643503..0000000 --- a/tests/models/network/tool_run.py +++ /dev/null @@ -1,8 +0,0 @@ -from uuid import uuid4 - -from generalresearch.models.network.tool_run import new_tool_run_from_nmap - - -def test_new_tool_run_from_nmap(nmap_run): - scan_group_id = uuid4().hex - run, scan = new_tool_run_from_nmap(nmap_run, scan_group_id=scan_group_id) |