mtr manager and tool run + test

9 files changed, 441 insertions, 32 deletions

diff --git a/generalresearch/managers/network/mtr.py b/generalresearch/managers/network/mtr.py
new file mode 100644
index 0000000..496d3f0
--- /dev/null
+++ b/generalresearch/managers/network/mtr.py
@@ -0,0 +1,47 @@
+from typing import Optional
+
+from psycopg import Cursor
+
+from generalresearch.managers.base import PostgresManager
+from generalresearch.models.network.tool_run import MtrRun
+
+
+class MtrManager(PostgresManager):
+
+    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 = """
+        INSERT INTO network_mtr (
+            run_id, source_ip, facility_id,
+            protocol, port, parsed
+        )
+        VALUES (
+            %(run_id)s, %(source_ip)s, %(facility_id)s,
+            %(protocol)s, %(port)s, %(parsed)s
+        );
+        """
+        params = run.model_dump_postgres()
+
+        query_hops = """
+        INSERT INTO network_mtrhop (
+            hop, ip, domain, asn, mtr_run_id
+        ) VALUES (
+            %(hop)s, %(ip)s, %(domain)s,
+            %(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]
+
+        if c:
+            c.execute(query, params)
+            if params_hops:
+                c.executemany(query_hops, params_hops)
+        else:
+            with self.pg_config.make_connection() as conn:
+                with conn.cursor() as c:
+                    c.execute(query, params)
+                    if params_hops:
+                        c.executemany(query_hops, params_hops)
diff --git a/generalresearch/managers/network/tool_run.py b/generalresearch/managers/network/tool_run.py
index 75c2e73..6280221 100644
--- a/generalresearch/managers/network/tool_run.py
+++ b/generalresearch/managers/network/tool_run.py
@@ -4,9 +4,15 @@ 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
+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.pg_helper import PostgresConfig
 
 
@@ -19,8 +25,9 @@ class ToolRunManager(PostgresManager):
         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)
 
-    def create_tool_run(self, run: PortScanRun | RDnsRun, c: Cursor):
+    def create_tool_run(self, run: PortScanRun | RDnsRun | MtrRun, c: Cursor):
         query = sql.SQL(
             """
         INSERT INTO network_toolrun (
@@ -88,3 +95,21 @@ class ToolRunManager(PostgresManager):
         )
         res["parsed"] = parsed
         return RDnsRun.model_validate(res)
+
+    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:
+        query = """
+        SELECT tr.*, mtr.parsed, mtr.source_ip, mtr.facility_id
+        FROM network_toolrun tr
+        JOIN network_mtr mtr ON tr.id = mtr.run_id
+        WHERE id = %(id)s
+        """
+        params = {"id": id}
+        res = self.pg_config.execute_sql_query(query, params)[0]
+        return MtrRun.model_validate(res)
diff --git a/generalresearch/models/network/mtr.py b/generalresearch/models/network/mtr.py
index 2e994d4..4b040de 100644
--- a/generalresearch/models/network/mtr.py
+++ b/generalresearch/models/network/mtr.py
@@ -6,7 +6,14 @@ from ipaddress import ip_address
 from typing import List, Optional, Dict
 
 import tldextract
-from pydantic import Field, field_validator, BaseModel, ConfigDict, model_validator
+from pydantic import (
+    Field,
+    field_validator,
+    BaseModel,
+    ConfigDict,
+    model_validator,
+    computed_field,
+)
 
 from generalresearch.models.network.definitions import IPProtocol, get_ip_kind, IPKind
 
@@ -16,7 +23,7 @@ class MTRHop(BaseModel):
 
     hop: int = Field(alias="count")
     host: str
-    asn: Optional[str] = Field(default=None, alias="ASN")
+    asn: Optional[int] = Field(default=None, alias="ASN")
 
     loss_pct: float = Field(alias="Loss%")
     sent: int = Field(alias="Snt")
@@ -27,15 +34,19 @@ class MTRHop(BaseModel):
     worst_ms: float = Field(alias="Wrst")
     stdev_ms: float = Field(alias="StDev")
 
-    hostname: Optional[str] = None
+    hostname: Optional[str] = Field(
+        default=None, examples=["fixed-187-191-8-145.totalplay.net"]
+    )
     ip: Optional[str] = None
 
-    @field_validator("asn")
+    @field_validator("asn", mode="before")
     @classmethod
-    def normalize_asn(cls, v):
-        if v == "AS???":
+    def normalize_asn(cls, v: str):
+        if v is None or v == "AS???":
             return None
-        return v
+        if type(v) is int:
+            return v
+        return int(v.replace("AS", ""))
 
     @model_validator(mode="after")
     def parse_host(self):
@@ -72,11 +83,27 @@ class MTRHop(BaseModel):
             return False
         return self.stdev_ms > self.avg_ms or self.worst_ms > self.best_ms * 10
 
+    @computed_field(examples=["totalplay.net"])
     @cached_property
     def domain(self) -> Optional[str]:
         if self.hostname:
             return tldextract.extract(self.hostname).top_domain_under_public_suffix
 
+    def model_dump_postgres(self, run_id: int):
+        # Writes for the network_mtrhop table
+        d = {"mtr_run_id": run_id}
+        data = self.model_dump(
+            mode="json",
+            include={
+                "hop",
+                "ip",
+                "domain",
+                "asn",
+            },
+        )
+        d.update(data)
+        return d
+
 
 class MTRReport(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
@@ -97,8 +124,20 @@ class MTRReport(BaseModel):
 
     hops: List[MTRHop] = Field()
 
+    def model_dump_postgres(self):
+        # Writes for the network_mtr table
+        d = self.model_dump(
+            mode="json",
+            include={"port"},
+        )
+        d["protocol"] = self.protocol.to_number()
+        d["parsed"] = self.model_dump_json(indent=0)
+        return d
+
     def print_report(self) -> None:
-        print(f"MTR Report → {self.destination} {self.protocol.name} {self.port or ''}\n")
+        print(
+            f"MTR Report → {self.destination} {self.protocol.name} {self.port or ''}\n"
+        )
         host_max_len = max(len(h.host) for h in self.hops)
 
         header = (
@@ -198,8 +237,8 @@ def run_mtr(
     )
     raw = proc.stdout.strip()
     data = parse_raw_output(raw)
-    data['port'] = port
-    data['protocol'] = protocol
+    data["port"] = port
+    data["protocol"] = protocol
     return MTRReport.model_validate(data)
 
 
@@ -208,14 +247,3 @@ def parse_raw_output(raw: str) -> Dict:
     data.update(data.pop("mtr"))
     data["hops"] = data.pop("hubs")
     return data
-
-
-def load_example():
-    s = open(
-        "/home/gstupp/projects/generalresearch/generalresearch/models/network/mtr_fatbeam.json",
-        "r",
-    ).read()
-    data = parse_raw_output(s)
-    data['port'] = 443
-    data['protocol'] = IPProtocol.TCP
-    return MTRReport.model_validate(data)
diff --git a/generalresearch/models/network/rdns.py b/generalresearch/models/network/rdns.py
index ac63414..e00a32d 100644
--- a/generalresearch/models/network/rdns.py
+++ b/generalresearch/models/network/rdns.py
@@ -40,7 +40,7 @@ class RDNSResult(BaseModel):
     def hostname_count(self) -> int:
         return len(self.hostnames)
 
-    @computed_field(examples=["totalplay"])
+    @computed_field(examples=["totalplay.net"])
     @cached_property
     def primary_domain(self) -> Optional[str]:
         if self.primary_hostname:
diff --git a/generalresearch/models/network/tool_run.py b/generalresearch/models/network/tool_run.py
index fba5dcb..2588890 100644
--- a/generalresearch/models/network/tool_run.py
+++ b/generalresearch/models/network/tool_run.py
@@ -17,6 +17,12 @@ from generalresearch.models.network.rdns import (
     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
 
 
@@ -90,6 +96,20 @@ class RDnsRun(ToolRun):
         return d
 
 
+class MtrRun(ToolRun):
+    facility_id: int = Field(default=1)
+    source_ip: IPvAnyAddressStr = Field()
+    parsed: MTRReport = Field()
+
+    def model_dump_postgres(self):
+        d = super().model_dump_postgres()
+        d["run_id"] = self.id
+        d["source_ip"] = self.source_ip
+        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:
@@ -129,3 +149,25 @@ def run_dig(ip: str, scan_group_id: Optional[UUIDStr] = None) -> RDnsRun:
         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/thl_django/network/models.py b/generalresearch/thl_django/network/models.py
index d50a7b1..7d4d8de 100644
--- a/generalresearch/thl_django/network/models.py
+++ b/generalresearch/thl_django/network/models.py
@@ -208,6 +208,9 @@ class MTR(models.Model):
     # nullable b/c ICMP doesn't use ports
     port = models.PositiveIntegerField(null=True)
 
+    # Full parsed output
+    parsed = models.JSONField()
+
     class Meta:
         db_table = "network_mtr"
 
@@ -219,9 +222,8 @@ class MTRHop(models.Model):
         related_name="hops",
     )
 
-    hop_number = models.PositiveSmallIntegerField()
-
-    responder_ip = models.GenericIPAddressField(null=True)
+    hop = models.PositiveSmallIntegerField()
+    ip = models.GenericIPAddressField(null=True)
 
     domain = models.CharField(max_length=50, null=True)
     asn = models.PositiveIntegerField(null=True)
@@ -230,13 +232,12 @@ class MTRHop(models.Model):
         db_table = "network_mtrhop"
         constraints = [
             models.UniqueConstraint(
-                fields=["mtr_run", "hop_number"],
+                fields=["mtr_run", "hop"],
                 name="unique_hop_per_run",
             )
         ]
         indexes = [
-            models.Index(fields=["mtr_run", "hop_number"]),
-            models.Index(fields=["responder_ip"]),
+            models.Index(fields=["ip"]),
             models.Index(fields=["asn"]),
             models.Index(fields=["domain"]),
         ]
diff --git a/test_utils/conftest.py b/test_utils/conftest.py
index 54fb682..187ff58 100644
--- a/test_utils/conftest.py
+++ b/test_utils/conftest.py
@@ -38,6 +38,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)
 
diff --git a/tests/data/mtr_fatbeam.json b/tests/data/mtr_fatbeam.json
new file mode 100644
index 0000000..6e27eb1
--- /dev/null
+++ b/tests/data/mtr_fatbeam.json
@@ -0,0 +1,206 @@
+{
+    "report": {
+        "mtr": {
+            "src": "gstupp-ThinkPad-X1-Carbon-Gen-11",
+            "dst": "167.150.6.80",
+            "tos": 0,
+            "tests": 10,
+            "psize": "64",
+            "bitpattern": "0x00"
+        },
+        "hubs": [
+            {
+                "count": 1,
+                "host": "_gateway (172.20.20.1)",
+                "ASN": "AS???",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 2.408,
+                "Avg": 16.157,
+                "Best": 2.408,
+                "Wrst": 69.531,
+                "StDev": 20.69
+            },
+            {
+                "count": 2,
+                "host": "172.16.20.1",
+                "ASN": "AS???",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 3.411,
+                "Avg": 16.906,
+                "Best": 2.613,
+                "Wrst": 90.7,
+                "StDev": 27.547
+            },
+            {
+                "count": 3,
+                "host": "192.168.1.254",
+                "ASN": "AS???",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 17.012,
+                "Avg": 9.812,
+                "Best": 3.061,
+                "Wrst": 25.728,
+                "StDev": 8.908
+            },
+            {
+                "count": 4,
+                "host": "ipdsl-jal-ptovallarta-19-l0.uninet.net.mx (201.154.95.117)",
+                "ASN": "AS???",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 6.954,
+                "Avg": 10.216,
+                "Best": 6.177,
+                "Wrst": 16.151,
+                "StDev": 3.343
+            },
+            {
+                "count": 5,
+                "host": "bb-la-onewilshire-29-ae32_0.uninet.net.mx (189.246.202.49)",
+                "ASN": "AS???",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 52.557,
+                "Avg": 54.174,
+                "Best": 45.681,
+                "Wrst": 71.387,
+                "StDev": 8.011
+            },
+            {
+                "count": 6,
+                "host": "ae91.edge7.LosAngeles1.Level3.net (4.7.28.197)",
+                "ASN": "AS3356",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 1079.2,
+                "Avg": 875.97,
+                "Best": 47.78,
+                "Wrst": 4150.5,
+                "StDev": 1345.7
+            },
+            {
+                "count": 7,
+                "host": "???",
+                "ASN": "AS???",
+                "Loss%": 100.0,
+                "Snt": 10,
+                "Last": 0.0,
+                "Avg": 0.0,
+                "Best": 0.0,
+                "Wrst": 0.0,
+                "StDev": 0.0
+            },
+            {
+                "count": 8,
+                "host": "ae10.cr1.lax10.us.zip.zayo.com (64.125.28.224)",
+                "ASN": "AS6461",
+                "Loss%": 70.0,
+                "Snt": 10,
+                "Last": 1186.5,
+                "Avg": 2189.8,
+                "Best": 1186.5,
+                "Wrst": 3202.8,
+                "StDev": 1008.2
+            },
+            {
+                "count": 9,
+                "host": "ae16.cr1.sjc1.us.zip.zayo.com (64.125.21.171)",
+                "ASN": "AS6461",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 92.819,
+                "Avg": 414.75,
+                "Best": 90.799,
+                "Wrst": 2140.8,
+                "StDev": 690.96
+            },
+            {
+                "count": 10,
+                "host": "ae27.cs3.sjc7.us.zip.zayo.com (64.125.18.28)",
+                "ASN": "AS6461",
+                "Loss%": 90.0,
+                "Snt": 10,
+                "Last": 5234.8,
+                "Avg": 5234.8,
+                "Best": 5234.8,
+                "Wrst": 5234.8,
+                "StDev": 0.0
+            },
+            {
+                "count": 11,
+                "host": "???",
+                "ASN": "AS???",
+                "Loss%": 100.0,
+                "Snt": 10,
+                "Last": 0.0,
+                "Avg": 0.0,
+                "Best": 0.0,
+                "Wrst": 0.0,
+                "StDev": 0.0
+            },
+            {
+                "count": 12,
+                "host": "ae8.cr1.sea1.us.zip.zayo.com (64.125.28.193)",
+                "ASN": "AS6461",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 93.389,
+                "Avg": 1238.6,
+                "Best": 91.537,
+                "Wrst": 5223.9,
+                "StDev": 1644.1
+            },
+            {
+                "count": 13,
+                "host": "ae7.ter2.sea1.us.zip.zayo.com (64.125.19.197)",
+                "ASN": "AS6461",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 91.212,
+                "Avg": 112.17,
+                "Best": 90.979,
+                "Wrst": 178.26,
+                "StDev": 30.086
+            },
+            {
+                "count": 14,
+                "host": "208.185.33.178.IDIA-369396-ZYO.zip.zayo.com (208.185.33.178)",
+                "ASN": "AS6461",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 103.95,
+                "Avg": 104.46,
+                "Best": 90.349,
+                "Wrst": 136.62,
+                "StDev": 15.726
+            },
+            {
+                "count": 15,
+                "host": "168.245.215.250",
+                "ASN": "AS55039",
+                "Loss%": 0.0,
+                "Snt": 10,
+                "Last": 85.672,
+                "Avg": 95.289,
+                "Best": 84.352,
+                "Wrst": 156.16,
+                "StDev": 21.621
+            },
+            {
+                "count": 16,
+                "host": "???",
+                "ASN": "AS???",
+                "Loss%": 100.0,
+                "Snt": 10,
+                "Last": 0.0,
+                "Avg": 0.0,
+                "Best": 0.0,
+                "Wrst": 0.0,
+                "StDev": 0.0
+            }
+        ]
+    }
+}
diff --git a/tests/managers/network/tool_run.py b/tests/managers/network/tool_run.py
index a598a71..0f9388f 100644
--- a/tests/managers/network/tool_run.py
+++ b/tests/managers/network/tool_run.py
@@ -1,11 +1,27 @@
+import os
+from datetime import datetime, timezone
 from uuid import uuid4
 
 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()
 
 
@@ -38,6 +54,7 @@ def test_run_dig(toolrun_manager):
 
     assert reverse_dns_run == run_out
 
+
 def test_run_dig_empty(toolrun_manager):
     reverse_dns_run = run_dig(ip=fake.ipv6())
 
@@ -45,4 +62,46 @@ def test_run_dig_empty(toolrun_manager):
 
     run_out = toolrun_manager.get_rdns_run(reverse_dns_run.id)
 
-    assert reverse_dns_run == run_out
\ No newline at end of file
+    assert reverse_dns_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_mtr(toolrun_manager, mtr_report):
+    started_at = datetime.now(tz=timezone.utc)
+    tool_version = get_mtr_version()
+
+    ip = mtr_report.destination
+
+    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"
+    )
+
+    toolrun_manager.create_mtr_run(run)
+
+    run_out = toolrun_manager.get_mtr_run(run.id)
+
+    assert run == run_out