MTR: domain, protocol, port. Change django models for mtr instead of traceroute

4 files changed, 45 insertions, 74 deletions

diff --git a/generalresearch/managers/network/rdns.py b/generalresearch/managers/network/rdns.py
index 2eed303..0b9b7b6 100644
--- a/generalresearch/managers/network/rdns.py
+++ b/generalresearch/managers/network/rdns.py
@@ -14,11 +14,11 @@ class RdnsManager(PostgresManager):
         """
         query = """
         INSERT INTO network_rdnsresult (
-            run_id, primary_hostname, primary_org,
+            run_id, primary_hostname, primary_domain,
             hostname_count, hostnames
         )
         VALUES (
-            %(run_id)s, %(primary_hostname)s, %(primary_org)s,
+            %(run_id)s, %(primary_hostname)s, %(primary_domain)s,
             %(hostname_count)s, %(hostnames)s
         );
         """
diff --git a/generalresearch/models/network/mtr.py b/generalresearch/models/network/mtr.py
index 98e7c16..2e994d4 100644
--- a/generalresearch/models/network/mtr.py
+++ b/generalresearch/models/network/mtr.py
@@ -1,9 +1,11 @@
 import json
 import re
 import subprocess
+from functools import cached_property
 from ipaddress import ip_address
 from typing import List, Optional, Dict
 
+import tldextract
 from pydantic import Field, field_validator, BaseModel, ConfigDict, model_validator
 
 from generalresearch.models.network.definitions import IPProtocol, get_ip_kind, IPKind
@@ -60,16 +62,21 @@ class MTRHop(BaseModel):
         self.ip = None
         return self
 
-    @property
+    @cached_property
     def ip_kind(self) -> Optional[IPKind]:
         return get_ip_kind(self.ip)
 
-    @property
+    @cached_property
     def icmp_rate_limited(self):
         if self.avg_ms == 0:
             return False
         return self.stdev_ms > self.avg_ms or self.worst_ms > self.best_ms * 10
 
+    @cached_property
+    def domain(self) -> Optional[str]:
+        if self.hostname:
+            return tldextract.extract(self.hostname).top_domain_under_public_suffix
+
 
 class MTRReport(BaseModel):
     model_config = ConfigDict(populate_by_name=True)
@@ -83,10 +90,15 @@ class MTRReport(BaseModel):
     psize: int = Field(description="Probe packet size in bytes.")
     bitpattern: str = Field(description="Payload byte pattern used in probes (hex).")
 
+    # Protocol used for the traceroute
+    protocol: IPProtocol = Field()
+    # The target port number for TCP/SCTP/UDP traces
+    port: Optional[int] = Field()
+
     hops: List[MTRHop] = Field()
 
     def print_report(self) -> None:
-        print(f"MTR Report → {self.destination}\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 = (
@@ -186,6 +198,8 @@ def run_mtr(
     )
     raw = proc.stdout.strip()
     data = parse_raw_output(raw)
+    data['port'] = port
+    data['protocol'] = protocol
     return MTRReport.model_validate(data)
 
 
@@ -202,4 +216,6 @@ def load_example():
         "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 44697c7..ac63414 100644
--- a/generalresearch/models/network/rdns.py
+++ b/generalresearch/models/network/rdns.py
@@ -26,7 +26,7 @@ class RDNSResult(BaseModel):
         assert len(self.hostnames) == self.hostname_count
         if self.hostnames:
             assert self.hostnames[0] == self.primary_hostname
-            assert self.primary_org in self.primary_hostname
+            assert self.primary_domain in self.primary_hostname
         return self
 
     @computed_field(examples=["fixed-187-191-8-145.totalplay.net"])
@@ -42,15 +42,15 @@ class RDNSResult(BaseModel):
 
     @computed_field(examples=["totalplay"])
     @cached_property
-    def primary_org(self) -> Optional[str]:
+    def primary_domain(self) -> Optional[str]:
         if self.primary_hostname:
-            return tldextract.extract(self.primary_hostname).domain
+            return tldextract.extract(self.primary_hostname).top_domain_under_public_suffix
 
     def model_dump_postgres(self):
         # Writes for the network_rdnsresult table
         d = self.model_dump(
             mode="json",
-            include={"primary_hostname", "primary_org", "hostname_count"},
+            include={"primary_hostname", "primary_domain", "hostname_count"},
         )
         d["hostnames"] = json.dumps(self.hostnames)
         return d
diff --git a/generalresearch/thl_django/network/models.py b/generalresearch/thl_django/network/models.py
index b0f4cdc..d50a7b1 100644
--- a/generalresearch/thl_django/network/models.py
+++ b/generalresearch/thl_django/network/models.py
@@ -98,7 +98,7 @@ class RDNSResult(models.Model):
     )
 
     primary_hostname = models.CharField(max_length=255, null=True)
-    primary_org = models.CharField(max_length=50, null=True)
+    primary_domain = models.CharField(max_length=50, null=True)
     hostname_count = models.PositiveIntegerField(default=0)
     hostnames = models.JSONField(default=list)
 
@@ -106,7 +106,7 @@ class RDNSResult(models.Model):
         db_table = "network_rdnsresult"
         indexes = [
             models.Index(fields=["primary_hostname"]),
-            models.Index(fields=["primary_org"]),
+            models.Index(fields=["primary_domain"]),
         ]
 
 
@@ -191,11 +191,11 @@ class PortScanPort(models.Model):
         ]
 
 
-class Traceroute(models.Model):
+class MTR(models.Model):
     run = models.OneToOneField(
         ToolRun,
         on_delete=models.CASCADE,
-        related_name="traceroute",
+        related_name="mtr",
         primary_key=True,
     )
 
@@ -204,88 +204,42 @@ class Traceroute(models.Model):
     facility_id = models.PositiveIntegerField()
 
     # IANA protocol numbers (1=ICMP, 6=TCP, 17=UDP)
-    protocol = models.PositiveSmallIntegerField(default=1)
-
-    max_hops = models.PositiveSmallIntegerField()
-
-    # High-level result summary
-    final_responded = models.BooleanField()
-    reached_hop = models.PositiveSmallIntegerField(null=True)
-    total_duration_ms = models.PositiveIntegerField(null=True)
+    protocol = models.PositiveSmallIntegerField()
+    # nullable b/c ICMP doesn't use ports
+    port = models.PositiveIntegerField(null=True)
 
     class Meta:
-        db_table = "network_traceroute"
+        db_table = "network_mtr"
 
 
-class TracerouteHop(models.Model):
-    traceroute = models.ForeignKey(
-        Traceroute,
+class MTRHop(models.Model):
+    mtr_run = models.ForeignKey(
+        MTR,
         on_delete=models.CASCADE,
         related_name="hops",
     )
 
     hop_number = models.PositiveSmallIntegerField()
-    probe_number = models.PositiveSmallIntegerField()
 
     responder_ip = models.GenericIPAddressField(null=True)
 
-    rtt_ms = models.FloatField(null=True)
-
-    ttl = models.PositiveSmallIntegerField(null=True)
-
-    icmp_type = models.PositiveSmallIntegerField(null=True)
-    icmp_code = models.PositiveSmallIntegerField(null=True)
+    domain = models.CharField(max_length=50, null=True)
+    asn = models.PositiveIntegerField(null=True)
 
     class Meta:
-        db_table = "network_traceroutehop"
+        db_table = "network_mtrhop"
         constraints = [
             models.UniqueConstraint(
-                fields=["traceroute", "hop_number", "probe_number"],
-                name="unique_probe_per_hop",
+                fields=["mtr_run", "hop_number"],
+                name="unique_hop_per_run",
             )
         ]
         indexes = [
-            models.Index(fields=["traceroute", "hop_number"]),
+            models.Index(fields=["mtr_run", "hop_number"]),
             models.Index(fields=["responder_ip"]),
+            models.Index(fields=["asn"]),
+            models.Index(fields=["domain"]),
         ]
-        ordering = ["traceroute_id", "hop_number", "probe_number"]
-
-    def __str__(self):
-        return f"{self.traceroute} hop {self.hop_number}.{self.probe_number}"
-
-
-# class TracerouteAnalysis(models.Model):
-#     traceroute = models.OneToOneField(
-#         Traceroute,
-#         on_delete=models.CASCADE,
-#         related_name="analysis",
-#         primary_key=True,
-#     )
-#
-#     reached_destination = models.BooleanField()
-#
-#     hop_count = models.PositiveSmallIntegerField()
-#
-#     latency_spike_detected = models.BooleanField(default=False)
-#
-#     max_rtt_ms = models.FloatField(null=True)
-#     rtt_stddev = models.FloatField(null=True)
-#
-#     last_hop_private = models.BooleanField(default=False)
-#     last_hop_asn = models.PositiveIntegerField(null=True)
-#
-#     # Deterministic hash of first N hops (binary SHA256 recommended)
-#     path_prefix_hash = models.BinaryField(max_length=32, null=True)
-#
-#     anomaly_score = models.FloatField(null=True)
-#
-#     class Meta:
-#         db_table = "network_tracerouteanalysis"
-#         indexes = [
-#             models.Index(fields=["path_prefix_hash"]),
-#             models.Index(fields=["anomaly_score"]),
-#         ]
-#
 
 
 class IPLabel(models.Model):
@@ -293,6 +247,7 @@ class IPLabel(models.Model):
     Stores *ground truth* about an IP at a specific time.
     Used for model training and evaluation.
     """
+
     id = models.BigAutoField(primary_key=True, null=False)
 
     ip = CIDRField()