From a13fc5b282f3f294ababf0afcfdfda3c28a86662 Mon Sep 17 00:00:00 2001
From: Kris Forbes <kris@gohosted.ca>
Date: Fri, 6 Feb 2026 16:59:56 -0500
Subject: [PATCH] Strip .prod.global.gc.ca from hostnames

---
 test_inventory_keys.py   | 7 +++++++
 wif2ansible/inventory.py | 8 +++++++-
 wif2ansible/parsers.py   | 9 +++++++++
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/test_inventory_keys.py b/test_inventory_keys.py
index d79256c..eacd127 100644
--- a/test_inventory_keys.py
+++ b/test_inventory_keys.py
@@ -55,5 +55,12 @@ class TestInventoryKeys(unittest.TestCase):
         self.assertIn("good_name", hosts)
         self.assertNotIn("bad_name", hosts)
 
+    def test_suffix_stripping(self):
+        from wif2ansible.parsers import clean_hostname
+        self.assertEqual(clean_hostname("server.prod.global.gc.ca"), "server")
+        self.assertEqual(clean_hostname("server.PROD.GLOBAL.GC.CA"), "server")
+        self.assertEqual(clean_hostname("nosuffix"), "nosuffix")
+        self.assertEqual(clean_hostname("other.suffix.com"), "other.suffix.com")
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/wif2ansible/inventory.py b/wif2ansible/inventory.py
index 75c9de3..a537e0a 100644
--- a/wif2ansible/inventory.py
+++ b/wif2ansible/inventory.py
@@ -117,10 +117,16 @@ def generate_inventory(servers: Dict[str, Server], flows: List[Flow]) -> Dict[st
                  print(f"Warning: No resolvable name found for {server.primary_ip} (Candidates: {candidates}). Using IP.")
             final_host_key = server.primary_ip
             
-        host_key = final_host_key
+        # Final cleanup: Strip suffixes if user requested
+        from .parsers import clean_hostname
+        host_key = clean_hostname(final_host_key)
         
         if host_key not in inventory_hosts:
             host_vars = server.get_ansible_vars()
+            
+            # Ensure proper ansible_host is set if key is not IP
+            if host_key != server.primary_ip and server.primary_ip:
+                 host_vars['ansible_host'] = server.primary_ip
             host_vars['flows'] = []
             inventory_hosts[host_key] = host_vars
             
diff --git a/wif2ansible/parsers.py b/wif2ansible/parsers.py
index 991b8a3..f02b32c 100644
--- a/wif2ansible/parsers.py
+++ b/wif2ansible/parsers.py
@@ -72,6 +72,15 @@ def clean_reference(ref: str) -> str:
     # Remove leading/trailing whitespace
     return s.strip()
 
+def clean_hostname(name: str) -> str:
+    """
+    Strips specific suffixes like .prod.global.gc.ca to get shortname.
+    """
+    if not name:
+        return ""
+    # Case insensitive strip
+    return re.sub(r'\.prod\.global\.gc\.ca$', '', name, flags=re.IGNORECASE)
+
 def parse_ip(ip_str: str) -> List[str]:
     """Finds all IPv4 addresses in a string."""
     if not ip_str: