From 9e28004d6cd4d9825e0af2cd0b7a92796a50f352 Mon Sep 17 00:00:00 2001
From: Kris Forbes <kris@gohosted.ca>
Date: Fri, 6 Feb 2026 15:56:58 -0500
Subject: [PATCH] Strip SRV prefix from inventory host keys

---
 test_inventory_keys.py      | 31 +++++++++++++++++++++++++++++++
 wif2ansible/excel_reader.py |  4 ++--
 wif2ansible/parsers.py      | 15 +++++++++++++++
 3 files changed, 48 insertions(+), 2 deletions(-)
 create mode 100644 test_inventory_keys.py

diff --git a/test_inventory_keys.py b/test_inventory_keys.py
new file mode 100644
index 0000000..2779bef
--- /dev/null
+++ b/test_inventory_keys.py
@@ -0,0 +1,31 @@
+import unittest
+from wif2ansible.models import Server, Flow
+from wif2ansible.inventory import generate_inventory
+
+class TestInventoryKeys(unittest.TestCase):
+    def test_inventory_keys_are_hostnames(self):
+        # Create a server with Ref, Hostname, IP
+        s1 = Server(reference="SERVER_REF_01", hostname="server01", ip_address="192.168.1.10", platform="windows")
+        
+        # Create a flow matching this server
+        f1 = Flow(flow_id="1", source_ip="192.168.1.10", destination_ip="10.0.0.1", ports=[80])
+        
+        servers = {"SERVER_REF_01": s1}
+        flows = [f1]
+        
+        inventory = generate_inventory(servers, flows)
+        
+        # Verify stricture
+        hosts = inventory['all']['hosts']
+        
+        # Key should be REFERENCE "SERVER_REF_01" (or hostname/ip fallback)
+        self.assertIn("SERVER_REF_01", hosts)
+        self.assertNotIn("192.168.1.10", hosts)
+        
+        # Check variables
+        host_vars = hosts["SERVER_REF_01"]
+        self.assertEqual(host_vars['ansible_host'], "192.168.1.10")
+        self.assertEqual(host_vars['ansible_connection'], "winrm")
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/wif2ansible/excel_reader.py b/wif2ansible/excel_reader.py
index 3af0419..ced6bc4 100644
--- a/wif2ansible/excel_reader.py
+++ b/wif2ansible/excel_reader.py
@@ -2,7 +2,7 @@ import openpyxl
 from openpyxl.worksheet.worksheet import Worksheet
 from typing import List, Dict, Tuple, Optional
 from .models import Server, Flow
-from .parsers import parse_ports, parse_ip, clean_header
+from .parsers import parse_ports, parse_ip, clean_header, clean_reference
 
 from openpyxl.utils import get_column_letter
 
@@ -130,7 +130,7 @@ def read_servers(filename: str) -> Dict[str, Server]:
         
         s = Server(
             reference=ref, 
-            hostname=ref, 
+            hostname=clean_reference(ref), 
             platform=plat,
             ip_address=ip_addr,
             production_ip=prod_ip_addr
diff --git a/wif2ansible/parsers.py b/wif2ansible/parsers.py
index de8d2c5..a833816 100644
--- a/wif2ansible/parsers.py
+++ b/wif2ansible/parsers.py
@@ -57,6 +57,21 @@ def parse_ports(port_str: str) -> List[int]:
             
     return sorted(list(ports))
 
+def clean_reference(ref: str) -> str:
+    """
+    Cleans a server reference string.
+    Specifically removes 'SRV###' type prefixes if present.
+    Example: 'SRV123 MyServer' -> 'MyServer'
+    """
+    if not ref:
+        return ""
+        
+    s = str(ref)
+    # Remove SRV followed by digits and whitespace
+    s = re.sub(r'SRV\d+\s*', '', s, flags=re.IGNORECASE)
+    # Remove leading/trailing whitespace
+    return s.strip()
+
 def parse_ip(ip_str: str) -> List[str]:
     """Finds all IPv4 addresses in a string."""
     if not ip_str: