1 files changed, 20 insertions, 24 deletions

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)