Support for geored-based IP lookup on AAR

nickvsnetworking · Oct 23, 2024 · 51bfbdc · 51bfbdc
1 parent 030dc72
commit 51bfbdc
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 2 deletions.
diff --git a/lib/database.py b/lib/database.py
@@ -2190,7 +2190,7 @@ def Get_UE_by_IP(self, subscriber_routing):
         self.logTool.log(service='Database', level='debug', message="Called Get_UE_by_IP() for IP " + str(subscriber_routing), redisClient=self.redisMessaging)
 
         Session = sessionmaker(bind = self.engine)
-        session = Session()    
+        session = Session()
 
         try:
             result = session.query(SERVING_APN).filter_by(subscriber_routing=subscriber_routing).one()

diff --git a/lib/diameter.py b/lib/diameter.py
@@ -15,6 +15,7 @@
 import json
 import time
 import socket
+import requests
 import traceback
 import re
 from baseModels import Peer, OutboundData
@@ -2989,6 +2990,8 @@ def Answer_16777236_265(self, packet_vars, avps):
             msisdn = None
             identifier = None
             apnId = None
+            remoteServingApn = None
+            servingApn = None
             try:
                 serviceUrn = bytes.fromhex(self.get_avp_data(avps, 525)[0]).decode('ascii')
             except:
@@ -3025,6 +3028,18 @@ def Answer_16777236_265(self, packet_vars, avps):
                     if ipServingApn:
                         ipApnName = self.database.Get_APN(apn_id=int(ipServingApn.get('apn', {})))
                         ipApnName = ipApnName.get('apn', None)
+                    else:
+                        #If we didn't find a serving APN for the IP, try the other local HSS'.
+                        localGeoredEndpoints = self.config.get('geored', {}).get('local_endpoints', [])
+                        for localGeoredEndpoint in localGeoredEndpoints:
+                            response = requests.get(url=f"{localGeoredEndpoint}/pcrf/pcrf_serving_apn_ip/{ueIp}")
+                            responseJson = response.json()
+                            if not responseJson:
+                                continue
+                            remoteServingApn = responseJson
+                            ipApnName = self.database.Get_APN(apn_id=int(remoteServingApn.get('apn', {})))
+                            ipApnName = ipApnName.get('apn', None)
+
             except Exception as e:
                 self.logTool.log(service='HSS', level='debug', message=f"[diameter.py] [Answer_16777236_265] [AAA] Exception: {traceback.format_exc()}", redisClient=self.redisMessaging)
                 ipApnName = ''
@@ -3138,7 +3153,10 @@ def Answer_16777236_265(self, packet_vars, avps):
                                 apnId = (self.database.Get_APN_by_Name(apn="ims")).get('apn_id', None)
                                 self.logTool.log(service='HSS', level='debug', message=f"[diameter.py] [Answer_16777236_265] [AAA] ApnID: {apnId}", redisClient=self.redisMessaging)
                             self.logTool.log(service='HSS', level='debug', message=f"[diameter.py] [Answer_16777236_265] [AAA] Getting Serving APN for subscriberId: {subscriberId} and apnId: {apnId}", redisClient=self.redisMessaging)
-                            servingApn = self.database.Get_Serving_APN(subscriber_id=subscriberId, apn_id=apnId)
+                            if remoteServingApn:
+                                servingApn = remoteServingApn
+                            else:
+                                servingApn = self.database.Get_Serving_APN(subscriber_id=subscriberId, apn_id=apnId)
                             servingPgwPeer = servingApn.get('serving_pgw_peer', None).split(';')[0]
                             servingPgw = servingApn.get('serving_pgw', None)
                             servingPgwRealm = servingApn.get('serving_pgw_realm', None)

diff --git a/services/apiService.py b/services/apiService.py
@@ -1745,6 +1745,23 @@ def get(self, imsi, apn_id):
 
             return handle_exception(E)
 
+@ns_pcrf.route('/pcrf_serving_apn_ip/<string:ip_address>')
+class PyHSS_PCRF_Get_Serving_APN_IP(Resource):
+    def get(self, ip_address):
+        '''Get Serving APN Data for an IP Address'''
+        try:
+            serving_apn_final = {}
+            try:
+                serving_apn = databaseClient.Get_Serving_APN_By_IP(str(ip_address))
+            except:
+                serving_apn = None
+            if serving_apn:
+                serving_apn_final = databaseClient.Sanitize_Datetime(serving_apn)
+            return serving_apn_final, 200
+        except Exception as E:
+            print("Flask Exception: " + str(E))
+            return handle_exception(E)
+
 @ns_pcrf.route('/')
 class PyHSS_PCRF(Resource):
     @ns_pcrf.doc('Push Charging Rule to a Subscriber')