diff options
17 files changed, 187 insertions, 62 deletions
diff --git a/generalresearch/models/network/mtr/command.py b/generalresearch/models/network/mtr/command.py index e3ab903..f8d2d49 100644 --- a/generalresearch/models/network/mtr/command.py +++ b/generalresearch/models/network/mtr/command.py @@ -4,6 +4,7 @@ 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 +from generalresearch.models.network.tool_run_command import MTRRunCommand SUPPORTED_PROTOCOLS = { IPProtocol.TCP, @@ -56,20 +57,14 @@ def get_mtr_version() -> str: 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 - ) +def run_mtr(config: MTRRunCommand) -> MTRResult: + cmd = config.to_command_str() + args = cmd.split(" ") proc = subprocess.run( - args.split(" "), + args, capture_output=True, text=True, check=False, ) raw = proc.stdout.strip() - return parse_mtr_output(raw, protocol=protocol, port=port) + return parse_mtr_output(raw, protocol=config.options.protocol, port=config.options.port) diff --git a/generalresearch/models/network/mtr/execute.py b/generalresearch/models/network/mtr/execute.py index 81de24f..a6fb82a 100644 --- a/generalresearch/models/network/mtr/execute.py +++ b/generalresearch/models/network/mtr/execute.py @@ -10,30 +10,33 @@ from generalresearch.models.network.mtr.command import ( 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.tool_run_command import ( + MTRRunCommand, + MTRRunCommandOptions, +) from generalresearch.models.network.utils import get_source_ip def execute_mtr( ip: str, scan_group_id: Optional[UUIDStr] = None, - protocol: Optional[IPProtocol] = None, + protocol: Optional[IPProtocol] = IPProtocol.ICMP, port: Optional[int] = None, report_cycles: int = 10, ) -> MTRRun: + config = MTRRunCommand( + options=MTRRunCommandOptions( + ip=ip, + report_cycles=report_cycles, + protocol=protocol, + port=port, + ), + ) + started_at = datetime.now(tz=timezone.utc) tool_version = get_mtr_version() - result = run_mtr(ip, protocol=protocol, port=port, report_cycles=report_cycles) + result = run_mtr(config) finished_at = datetime.now(tz=timezone.utc) - raw_command = build_mtr_command(ip) - config = ToolRunCommand( - command="mtr", - options={ - "protocol": result.protocol, - "port": result.port, - "report_cycles": report_cycles, - }, - ) return MTRRun( tool_name=ToolName.MTR, @@ -43,7 +46,7 @@ def execute_mtr( ip=ip, started_at=started_at, finished_at=finished_at, - raw_command=raw_command, + raw_command=config.to_command_str(), scan_group_id=scan_group_id or uuid4().hex, config=config, parsed=result, diff --git a/generalresearch/models/network/mtr/parser.py b/generalresearch/models/network/mtr/parser.py index fb0ca61..685eeca 100644 --- a/generalresearch/models/network/mtr/parser.py +++ b/generalresearch/models/network/mtr/parser.py @@ -5,10 +5,10 @@ from generalresearch.models.network.definitions import IPProtocol from generalresearch.models.network.mtr.result import MTRResult -def parse_mtr_output(raw: str, port, protocol) -> MTRResult: +def parse_mtr_output(raw: str, port: int, protocol: IPProtocol) -> MTRResult: data = parse_mtr_raw_output(raw) data["port"] = port - data["protocol"] = protocol or IPProtocol.ICMP + data["protocol"] = protocol return MTRResult.model_validate(data) diff --git a/generalresearch/models/network/nmap/command.py b/generalresearch/models/network/nmap/command.py index dfa55de..47e0a87 100644 --- a/generalresearch/models/network/nmap/command.py +++ b/generalresearch/models/network/nmap/command.py @@ -3,18 +3,40 @@ from typing import Optional, List from generalresearch.models.network.nmap.parser import parse_nmap_xml from generalresearch.models.network.nmap.result import NmapResult +from generalresearch.models.network.tool_run_command import NmapRunCommand -def build_nmap_command(ip: str, top_ports: Optional[int] = 1000) -> List[str]: +def build_nmap_command( + ip: str, + no_ping: bool = True, + enable_advanced: bool = True, + timing: int = 4, + ports: Optional[str] = None, + top_ports: Optional[int] = None, +) -> 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 + args = ["nmap"] + assert 0 <= timing <= 5 + args.append(f"-T{timing}") + if no_ping: + args.append("-Pn") + if enable_advanced: + args.append("-A") + if ports is not None: + assert top_ports is None + args.extend(["-p", ports]) + if top_ports is not None: + assert ports is None + args.extend(["--top-ports", str(top_ports)]) + args.extend(["-oX", "-", ip]) + return " ".join(args) -def run_nmap(ip: str, top_ports: Optional[int] = 1000) -> NmapResult: - args = build_nmap_command(ip=ip, top_ports=top_ports) + +def run_nmap(config: NmapRunCommand) -> NmapResult: + cmd = config.to_command_str() + args = cmd.split(" ") proc = subprocess.run( args, capture_output=True, diff --git a/generalresearch/models/network/nmap/execute.py b/generalresearch/models/network/nmap/execute.py index 68a9926..0334f50 100644 --- a/generalresearch/models/network/nmap/execute.py +++ b/generalresearch/models/network/nmap/execute.py @@ -4,15 +4,35 @@ 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 +from generalresearch.models.network.tool_run_command import ( + NmapRunCommand, + NmapRunCommandOptions, +) def execute_nmap( - ip: str, top_ports: Optional[int] = 1000, scan_group_id: Optional[UUIDStr] = None + ip: str, + top_ports: Optional[int] = 1000, + ports: Optional[str] = None, + no_ping: bool = True, + enable_advanced: bool = True, + timing: int = 4, + scan_group_id: Optional[UUIDStr] = None, ): - result = run_nmap(ip=ip, top_ports=top_ports) + config = NmapRunCommand( + options=NmapRunCommandOptions( + top_ports=top_ports, + ports=ports, + no_ping=no_ping, + enable_advanced=enable_advanced, + timing=timing, + ip=ip, + ) + ) + result = run_nmap(config) assert result.exit_status == "success" assert result.target_ip == ip, f"{result.target_ip=}, {ip=}" + assert result.command_line == config.to_command_str() run = NmapRun( tool_name=ToolName.NMAP, @@ -24,7 +44,7 @@ def execute_nmap( 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}), + config=config, parsed=result, ) - return run
\ No newline at end of file + return run diff --git a/generalresearch/models/network/rdns/command.py b/generalresearch/models/network/rdns/command.py index e9d5bfd..aa48f2a 100644 --- a/generalresearch/models/network/rdns/command.py +++ b/generalresearch/models/network/rdns/command.py @@ -2,10 +2,12 @@ import subprocess from generalresearch.models.network.rdns.parser import parse_rdns_output from generalresearch.models.network.rdns.result import RDNSResult +from generalresearch.models.network.tool_run_command import RDNSRunCommand -def run_rdns(ip: str) -> RDNSResult: - args = build_rdns_command(ip).split(" ") +def run_rdns(config: RDNSRunCommand) -> RDNSResult: + cmd = config.to_command_str() + args = cmd.split(" ") proc = subprocess.run( args, capture_output=True, @@ -13,10 +15,10 @@ def run_rdns(ip: str) -> RDNSResult: check=False, ) raw = proc.stdout.strip() - return parse_rdns_output(ip, raw) + return parse_rdns_output(ip=config.options.ip, raw=raw) -def build_rdns_command(ip: str): +def build_rdns_command(ip: str) -> str: # e.g. dig +noall +answer -x 1.2.3.4 return " ".join(["dig", "+noall", "+answer", "-x", ip]) diff --git a/generalresearch/models/network/rdns/execute.py b/generalresearch/models/network/rdns/execute.py index 97b8bf8..03a5080 100644 --- a/generalresearch/models/network/rdns/execute.py +++ b/generalresearch/models/network/rdns/execute.py @@ -14,15 +14,18 @@ from generalresearch.models.network.tool_run import ( Status, RDNSRun, ) -from generalresearch.models.network.tool_run_command import ToolRunCommand +from generalresearch.models.network.tool_run_command import ( + RDNSRunCommand, + RDNSRunCommandOptions, +) 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) + config = RDNSRunCommand(options=RDNSRunCommandOptions(ip=ip)) + result = run_rdns(config) finished_at = datetime.now(tz=timezone.utc) - raw_command = build_rdns_command(ip) run = RDNSRun( tool_name=ToolName.DIG, @@ -32,9 +35,9 @@ def execute_rdns(ip: str, scan_group_id: Optional[UUIDStr] = None): ip=ip, started_at=started_at, finished_at=finished_at, - raw_command=raw_command, + raw_command=config.to_command_str(), scan_group_id=scan_group_id or uuid4().hex, - config=ToolRunCommand(command="dig", options={}), + config=config, parsed=result, ) diff --git a/generalresearch/models/network/rdns/parser.py b/generalresearch/models/network/rdns/parser.py index f12a6f4..231949e 100644 --- a/generalresearch/models/network/rdns/parser.py +++ b/generalresearch/models/network/rdns/parser.py @@ -2,12 +2,13 @@ import ipaddress import re from typing import List +from generalresearch.models.custom_types import IPvAnyAddressStr from generalresearch.models.network.rdns.result import RDNSResult PTR_RE = re.compile(r"\sPTR\s+([^\s]+)\.") -def parse_rdns_output(ip, raw): +def parse_rdns_output(ip: IPvAnyAddressStr, raw:str): hostnames: List[str] = [] for line in raw.splitlines(): diff --git a/generalresearch/models/network/tool_run.py b/generalresearch/models/network/tool_run.py index 36e6950..114d4b6 100644 --- a/generalresearch/models/network/tool_run.py +++ b/generalresearch/models/network/tool_run.py @@ -12,7 +12,12 @@ from generalresearch.models.custom_types import ( 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 +from generalresearch.models.network.tool_run_command import ( + ToolRunCommand, + NmapRunCommand, + RDNSRunCommand, + MTRRunCommand, +) class ToolClass(StrEnum): @@ -68,6 +73,7 @@ class ToolRun(BaseModel): class NmapRun(ToolRun): tool_class: Literal[ToolClass.PORT_SCAN] = Field(default=ToolClass.PORT_SCAN) tool_name: Literal[ToolName.NMAP] = Field(default=ToolName.NMAP) + config: NmapRunCommand = Field() parsed: NmapResult = Field() @@ -81,6 +87,7 @@ class NmapRun(ToolRun): class RDNSRun(ToolRun): tool_class: Literal[ToolClass.RDNS] = Field(default=ToolClass.RDNS) tool_name: Literal[ToolName.DIG] = Field(default=ToolName.DIG) + config: RDNSRunCommand = Field() parsed: RDNSResult = Field() @@ -94,6 +101,7 @@ class RDNSRun(ToolRun): class MTRRun(ToolRun): tool_class: Literal[ToolClass.TRACEROUTE] = Field(default=ToolClass.TRACEROUTE) tool_name: Literal[ToolName.MTR] = Field(default=ToolName.MTR) + config: MTRRunCommand = Field() facility_id: int = Field(default=1) source_ip: IPvAnyAddressStr = Field() diff --git a/generalresearch/models/network/tool_run_command.py b/generalresearch/models/network/tool_run_command.py index 5abe670..68d2070 100644 --- a/generalresearch/models/network/tool_run_command.py +++ b/generalresearch/models/network/tool_run_command.py @@ -1,9 +1,64 @@ -from typing import Dict, Optional +from typing import Dict, Optional, Literal from pydantic import BaseModel, Field +from generalresearch.models.custom_types import IPvAnyAddressStr +from generalresearch.models.network.definitions import IPProtocol + class ToolRunCommand(BaseModel): - # todo: expand with arguments specific for each tool command: str = Field() options: Dict[str, Optional[str | int]] = Field(default_factory=dict) + + +class NmapRunCommandOptions(BaseModel): + ip: IPvAnyAddressStr + top_ports: Optional[int] = Field(default=1000) + ports: Optional[str] = Field(default=None) + no_ping: bool = Field(default=True) + enable_advanced: bool = Field(default=True) + timing: int = Field(default=4) + + +class NmapRunCommand(ToolRunCommand): + command: Literal["nmap"] = Field(default="nmap") + options: NmapRunCommandOptions = Field() + + def to_command_str(self): + from generalresearch.models.network.nmap.command import build_nmap_command + + options = self.options + return build_nmap_command(**options.model_dump()) + + +class RDNSRunCommandOptions(BaseModel): + ip: IPvAnyAddressStr + + +class RDNSRunCommand(ToolRunCommand): + command: Literal["dig"] = Field(default="dig") + options: RDNSRunCommandOptions = Field() + + def to_command_str(self): + from generalresearch.models.network.rdns.command import build_rdns_command + + options = self.options + return build_rdns_command(**options.model_dump()) + + +class MTRRunCommandOptions(BaseModel): + ip: IPvAnyAddressStr = Field() + protocol: IPProtocol = Field(default=IPProtocol.ICMP) + port: Optional[int] = Field(default=None) + report_cycles: int = Field(default=10) + + +class MTRRunCommand(ToolRunCommand): + command: Literal["mtr"] = Field(default="mtr") + options: MTRRunCommandOptions = Field() + + def to_command_str(self): + from generalresearch.models.network.mtr.command import build_mtr_command + + options = self.options + return build_mtr_command(**options.model_dump()) diff --git a/test_utils/managers/network/conftest.py b/test_utils/managers/network/conftest.py index f6a4078..979dd63 100644 --- a/test_utils/managers/network/conftest.py +++ b/test_utils/managers/network/conftest.py @@ -7,13 +7,18 @@ import pytest from generalresearch.managers.network.label import IPLabelManager from generalresearch.managers.network.tool_run import ToolRunManager 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 +from generalresearch.models.network.tool_run_command import ( + MTRRunCommand, + MTRRunCommandOptions, + RDNSRunCommand, + NmapRunCommand, + NmapRunCommandOptions, + RDNSRunCommandOptions, +) @pytest.fixture(scope="session") @@ -51,15 +56,21 @@ def nmap_result(nmap_raw_output): @pytest.fixture(scope="session") def nmap_run(nmap_result, scan_group_id): r = nmap_result + config = NmapRunCommand( + command="nmap", + options=NmapRunCommandOptions( + ip=r.target_ip, ports="22-1000,11000,1100,3389,61232", top_ports=None + ), + ) 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, + raw_command=config.to_command_str(), scan_group_id=scan_group_id, - config=ToolRunCommand(command="nmap"), + config=config, parsed=r, ) @@ -79,15 +90,16 @@ def rdns_run(rdns_result, scan_group_id): r = rdns_result ip = "45.33.32.156" utc_now = datetime.now(tz=timezone.utc) + config = RDNSRunCommand(command="dig", options=RDNSRunCommandOptions(ip=ip)) return RDNSRun( tool_version="1.2.3", status=Status.SUCCESS, ip=ip, started_at=utc_now, finished_at=utc_now + timedelta(seconds=1), - raw_command=build_rdns_command(ip=ip), + raw_command=config.to_command_str(), scan_group_id=scan_group_id, - config=ToolRunCommand(command="dig"), + config=config, parsed=r, ) @@ -109,6 +121,12 @@ def mtr_result(mtr_raw_output): def mtr_run(mtr_result, scan_group_id): r = mtr_result utc_now = datetime.now(tz=timezone.utc) + config = MTRRunCommand( + command="mtr", + options=MTRRunCommandOptions( + ip=r.destination, protocol=IPProtocol.TCP, port=443 + ), + ) return MTRRun( tool_version="1.2.3", @@ -116,11 +134,9 @@ def mtr_run(mtr_result, scan_group_id): 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 - ), + raw_command=config.to_command_str(), scan_group_id=scan_group_id, - config=ToolRunCommand(command="mtr"), + config=config, parsed=r, facility_id=1, source_ip="1.2.3.4", diff --git a/tests/managers/network/label.py b/tests/managers/network/test_label.py index 5b9a790..5b9a790 100644 --- a/tests/managers/network/label.py +++ b/tests/managers/network/test_label.py diff --git a/tests/managers/network/tool_run.py b/tests/managers/network/test_tool_run.py index a815809..a815809 100644 --- a/tests/managers/network/tool_run.py +++ b/tests/managers/network/test_tool_run.py diff --git a/tests/models/network/mtr.py b/tests/models/network/test_mtr.py index 2965300..2965300 100644 --- a/tests/models/network/mtr.py +++ b/tests/models/network/test_mtr.py diff --git a/tests/models/network/nmap.py b/tests/models/network/test_nmap.py index f034bf0..0be98d4 100644 --- a/tests/models/network/nmap.py +++ b/tests/models/network/test_nmap.py @@ -17,7 +17,7 @@ def resolve(host): def test_execute_nmap_scanme(toolrun_manager): ip = resolve("scanme.nmap.org") - run = execute_nmap(ip=ip, top_ports=20) + run = execute_nmap(ip=ip, top_ports=None, ports="20-30", enable_advanced=False) assert run.tool_name == ToolName.NMAP assert run.tool_class == ToolClass.PORT_SCAN assert run.ip == ip diff --git a/tests/models/network/nmap_parser.py b/tests/models/network/test_nmap_parser.py index 96d7b37..96d7b37 100644 --- a/tests/models/network/nmap_parser.py +++ b/tests/models/network/test_nmap_parser.py diff --git a/tests/models/network/rdns.py b/tests/models/network/test_rdns.py index e56c494..e56c494 100644 --- a/tests/models/network/rdns.py +++ b/tests/models/network/test_rdns.py |