From ddb8f08dfc23f7145e505c9ea4c09a03724daecf Mon Sep 17 00:00:00 2001
From: SimeonOnSecurity <4913771+simeononsecurity@users.noreply.github.com>
Date: Sat, 19 Oct 2024 01:49:26 -0500
Subject: [PATCH] Update main.py

https://www.tech-invite.com/3m23/toc/tinv-3gpp-23-003_zq.html#e-ax-E
---
 main.py | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 68 insertions(+), 4 deletions(-)

diff --git a/main.py b/main.py
index 47df7ab..3479a87 100644
--- a/main.py
+++ b/main.py
@@ -9,6 +9,9 @@
 def get_mccmnc_json_path():
     """
     Dynamically locate the path to the mccmnc.json file based on where the mccmnc module is installed.
+
+    Returns:
+        str: Path to the mccmnc.json file.
     """
     mccmnc_dir = os.path.dirname(mccmnc.__file__)
     json_path = os.path.join(mccmnc_dir, 'mccmnc.json')
@@ -17,6 +20,12 @@ def get_mccmnc_json_path():
 def load_json_file(json_path):
     """
     Load data from the specified JSON file.
+
+    Args:
+        json_path (str): Path to the JSON file.
+
+    Returns:
+        dict: Loaded JSON data.
     """
     if os.path.exists(json_path):
         with open(json_path, 'r', encoding='utf-8') as file:
@@ -27,28 +36,62 @@ def load_json_file(json_path):
 def save_json_file(data, json_path):
     """
     Save data to the specified JSON file.
+
+    Args:
+        data (dict): Data to be saved.
+        json_path (str): Path to the JSON file.
     """
     with open(json_path, 'w', encoding='utf-8') as file:
         json.dump(data, file, indent=4)
     print(f"JSON data updated in {json_path}")
 
-def construct_realm_url(mcc, mnc, use_pub=False):
+def construct_realm_url(mcc, mnc, nid=None, service_id=None, use_pub=False):
     """
-    Construct the realm URL using MCC and MNC values.
+    Construct the realm URL using MCC, MNC, NID, and service ID values.
     If use_pub is True, use 'pub.3gppnetwork.org' instead of '3gppnetwork.org'.
+
+    Args:
+        mcc (int): Mobile Country Code (3-digit).
+        mnc (int): Mobile Network Code (2 or 3-digit).
+        nid (str, optional): NID identifying a Stand-alone Non-Public Network (SNPN).
+        service_id (str, optional): Service ID that describes the service or operation.
+        use_pub (bool, optional): Whether to use 'pub.3gppnetwork.org' instead of '3gppnetwork.org'.
+
+    Returns:
+        str: Constructed realm URL.
     """
     realm_suffix = "pub.3gppnetwork.org" if use_pub else "3gppnetwork.org"
-    return f"wlan.mnc{mnc:03d}.mcc{mcc:03d}.{realm_suffix}"
+    mnc_str = f"{mnc:03d}"
+    mcc_str = f"{mcc:03d}"
+
+    if nid:
+        realm = f"wlan.nid{nid}.mnc{mnc_str}.mcc{mcc_str}.{realm_suffix}"
+    else:
+        realm = f"wlan.mnc{mnc_str}.mcc{mcc_str}.{realm_suffix}"
+
+    if service_id:
+        realm = f"{service_id}.{realm}"
+
+    return realm
 
 def validate_host(host):
     """
     Validate host format.
+
+    Args:
+        host (str): Hostname to be validated.
+
+    Returns:
+        bool: True if the host format is valid, False otherwise.
     """
     return all(c.isalnum() or c in '-._' for c in host)
 
 def setup_resolvers():
     """
     Setup DNS resolvers with custom DNS servers, rotating through them.
+
+    Returns:
+        list: List of configured DNS resolvers.
     """
     resolver_list = []
     dns_servers = ['1.1.1.1', '8.8.8.8', '9.9.9.9']
@@ -61,6 +104,13 @@ def setup_resolvers():
 def naptr_lookup(realm, resolver):
     """
     Perform NAPTR DNS lookup on the given realm using a specified resolver.
+
+    Args:
+        realm (str): The realm to perform the NAPTR lookup on.
+        resolver (dns.resolver.Resolver): The DNS resolver to use for the lookup.
+
+    Returns:
+        str: Replacement value from the NAPTR record if found, otherwise None.
     """
     try:
         print(f"Performing NAPTR lookup for {realm} using resolver {resolver.nameservers[0]}")
@@ -78,6 +128,13 @@ def naptr_lookup(realm, resolver):
 def srv_lookup(host, resolver):
     """
     Perform SRV DNS lookup on the given host using a specified resolver.
+
+    Args:
+        host (str): The hostname to perform the SRV lookup on.
+        resolver (dns.resolver.Resolver): The DNS resolver to use for the lookup.
+
+    Returns:
+        tuple: A tuple containing the target host and port if found, otherwise (None, None).
     """
     try:
         print(f"Performing SRV lookup for {host} using resolver {resolver.nameservers[0]}")
@@ -94,7 +151,14 @@ def srv_lookup(host, resolver):
 def check_realm_existence(mcc, mnc, resolvers):
     """
     Check if the realm exists by performing NAPTR and SRV lookups.
-    Returns True if a valid realm is found, along with the host and port.
+
+    Args:
+        mcc (int): Mobile Country Code (3-digit).
+        mnc (int): Mobile Network Code (2 or 3-digit).
+        resolvers (list): List of DNS resolvers to use for lookups.
+
+    Returns:
+        tuple: A tuple containing a boolean indicating success, the host, and the port.
     """
     # Construct initial realm URL for the .pub domain
     realm_url_pub = construct_realm_url(mcc, mnc, use_pub=True)