Merge pull request #14 from Ostorlab/fix_bus_refactor

refactor pseudo code and fix out selector
Ostorlab · Aug 30, 2022 · 18288a9 · 18288a9
2 parents c106c1f + 9b49fd2
commit 18288a9
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/agent/ip2geo_agent.py b/agent/ip2geo_agent.py
@@ -42,10 +42,12 @@ def process(self, message: m.Message) -> None:
         """
         logger.info('processing message of selector : %s', message.selector)
 
-        out_selector = f'{message.selector}.geolocation'
         ip = message.data['host']
-        mask = message.data.get('mask', '32')
-        network = ipaddress.ip_network(f'{ip}/{mask}')
+        mask = message.data.get('mask')
+        version = message.data.get('version')
+
+        out_selector = 'v3.asset.ip.v6.geolocation' if version == 6 else 'v3.asset.ip.v4.geolocation'
+        network = ipaddress.ip_network(f'{ip}/{mask}') if mask is not None else ipaddress.ip_network(f'{ip}')
 
         # classify ip range based on geolocation
         for result in self.ip_range_visitor.dichotomy_ip_network_visit(