catocli 2.0.1__py3-none-any.whl → 2.0.2__py3-none-any.whl

catocli/parsers/custom/export_rules/export_rules.py

@@ -3,7 +3,7 @@ import json
 import sys
 from graphql_client.api.call_api import ApiClient, CallApi
 from graphql_client.api_client import ApiException
-
+from ..customLib import writeDataToFile, makeCall, getAccountID
 
 def strip_ids_recursive(data):
     """Recursively strip id attributes from data structure, but only from objects that contain only 'id' and 'name' keys"""
@@ -37,40 +37,7 @@ def export_if_rules_to_json(args, configuration):
     Adapted from scripts/export_if_rules_to_json.py
     """
     try:
-        # Get account ID from args, configuration, or environment variable
-        account_id = None
-        if hasattr(args, 'accountID') and args.accountID:
-            account_id = args.accountID
-        elif hasattr(configuration, 'accountID') and configuration.accountID:
-            account_id = configuration.accountID
-        else:
-            account_id = os.getenv('CATO_ACCOUNT_ID')
-        
-        if not account_id:
-            raise ValueError("Account ID is required. Provide it using the -accountID flag or set CATO_ACCOUNT_ID environment variable.")
-        
-        # Set up output file path
-        if hasattr(args, 'output_file_path') and args.output_file_path:
-            # Use output file path if provided
-            output_file = args.output_file_path
-            destination_dir = os.path.dirname(output_file)
-            if hasattr(args, 'verbose') and args.verbose:
-                print(f"Using output file path: {output_file}")
-        else:
-            # Use default path and filename
-            destination_dir = 'config_data'
-            json_output_file = f"all_ifw_rules_and_sections_{account_id}.json"
-            output_file = os.path.join(destination_dir, json_output_file)
-            if hasattr(args, 'verbose') and args.verbose:
-                print(f"Using default path: {output_file}")
-        
-        # Create destination directory if it doesn't exist
-        if destination_dir and not os.path.exists(destination_dir):
-            if hasattr(args, 'verbose') and args.verbose:
-                print(f"Creating directory: {destination_dir}")
-            os.makedirs(destination_dir)
-        
-        # Define the GraphQL query
+        account_id = getAccountID(args, configuration)
         policy_query = {
             "query": "query policy ( $accountId:ID! ) { policy ( accountId:$accountId ) { internetFirewall { policy { enabled rules { audit { updatedTime updatedBy } rule { id name description index section { id name } enabled source { ip host { id name } site { id name } subnet ipRange { from to } globalIpRange { id name } networkInterface { id name } siteNetworkSubnet { id name } floatingSubnet { id name } user { id name } usersGroup { id name } group { id name } systemGroup { id name } } connectionOrigin country { id name } device { id name } deviceOS deviceAttributes { category type model manufacturer os osVersion } destination { application { id name } customApp { id name } appCategory { id name } customCategory { id name } sanctionedAppsCategory { id name } country { id name } domain fqdn ip subnet ipRange { from to } globalIpRange { id name } remoteAsn containers { fqdnContainer { id name } ipAddressRangeContainer { id name } } } service { standard { id name } custom { port portRange { from to } protocol } } action tracking { event { enabled } alert { enabled frequency subscriptionGroup { id name } webhook { id name } mailingList { id name } } } schedule { activeOn customTimeframePolicySchedule: customTimeframe { from to } customRecurringPolicySchedule: customRecurring { from to days } } exceptions { name source { ip host { id name } site { id name } subnet ipRange { from to } globalIpRange { id name } networkInterface { id name } siteNetworkSubnet { id name } floatingSubnet { id name } user { id name } usersGroup { id name } group { id name } systemGroup { id name } } deviceOS country { id name } device { id name } deviceAttributes { category type model manufacturer os osVersion } destination { application { id name } customApp { id name } appCategory { id name } customCategory { id name } sanctionedAppsCategory { id name } country { id name } domain fqdn ip subnet ipRange { from to } globalIpRange { id name } remoteAsn containers { fqdnContainer { id name } ipAddressRangeContainer { id name } } } service { standard { id name } custom { port portRangeCustomService: portRange { from to } protocol } } connectionOrigin } } properties } sections { audit { updatedTime updatedBy } section { id name } properties } audit { publishedTime publishedBy } revision { id name description changes createdTime updatedTime } } } } }",
             "variables": {
@@ -78,56 +45,7 @@ def export_if_rules_to_json(args, configuration):
             },
             "operationName": "policy"
         }
-        
-        # Create API client instance with params
-        instance = CallApi(ApiClient(configuration))
-        params = {
-            'v': hasattr(args, 'verbose') and args.verbose,  # verbose mode
-            'f': 'json',  # format
-            'p': False,  # pretty print
-            't': False   # test mode
-        }
-        
-        try:
-            # Call the API directly
-            # NOTE: The API client (graphql_client/api_client_types.py lines 106-108) 
-            # automatically prints error responses and exits on GraphQL errors.
-            # This means our custom error handling below may not be reached if there are GraphQL errors.
-            response = instance.call_api(policy_query, params)
-            all_ifw_rules = response[0] if response else {}
-            
-            # Show raw API response in verbose mode
-            if hasattr(args, 'verbose') and args.verbose:
-                print("\n" + "=" * 80)
-                print("RAW API RESPONSE:")
-                print("=" * 80)
-                print(json.dumps(all_ifw_rules, indent=2))
-                print("=" * 80 + "\n")
-            
-            # Check for GraphQL errors first (may not be reached due to API client behavior)
-            if 'errors' in all_ifw_rules:
-                error_messages = [error.get('message', 'Unknown error') for error in all_ifw_rules['errors']]
-                raise Exception(f"API returned errors: {', '.join(error_messages)}")
-            
-            if not all_ifw_rules or 'data' not in all_ifw_rules:
-                raise ValueError("Failed to retrieve data from API")
-                
-        except ApiException as e:
-            raise Exception(f"API call failed - {e}")
-        except Exception as e:
-            raise Exception(f"Unexpected error during API call - {e}")
-
-        # IMPORTANT: Preserve section IDs BEFORE stripping them
-        section_id_map = {}
-        section_to_start_after_id = None
-        sections_with_ids = all_ifw_rules['data']['policy']['internetFirewall']['policy']['sections']
-        for index, section_data in enumerate(sections_with_ids):
-            section_name = section_data['section']['name']
-            section_id = section_data['section']['id']
-            if index == 0:
-                section_to_start_after_id = section_id
-            else:
-                section_id_map[section_name] = section_id
+        all_ifw_rules = makeCall(args, configuration, policy_query)
             
         # Processing data to strip id attributes
         processed_data = strip_ids_recursive(all_ifw_rules)
@@ -173,26 +91,28 @@ def export_if_rules_to_json(args, configuration):
         
         # Reformat sections array to have index, section_id and section_name structure
         # Exclude the first section from export
+        sections_with_ids = all_ifw_rules['data']['policy']['internetFirewall']['policy']['sections']
         processed_sections = [] 
-        for index, section_data in enumerate(processed_data['data']['policy']['internetFirewall']['policy']['sections']):
+        for index, section_data in enumerate(sections_with_ids):
+            # print("sections_with_ids",json.dumps(section_data, indent=2))
             if index > 0:  # Skip the first section (index 0)
                 processed_sections.append({
                     "section_index": index,
-                    "section_name": section_data['section']['name']
+                    "section_name": section_data['section']['name'],
+                    "section_id": section_data['section']['id']
                 })
-
-        # Add preserved section IDs and section_to_start_after_id
-        processed_data['data']['policy']['internetFirewall']['policy']['section_ids'] = section_id_map
-        if section_to_start_after_id:
-            processed_data['data']['policy']['internetFirewall']['policy']['section_to_start_after_id'] = section_to_start_after_id
         
         # Replace the original sections array with the reformatted one
         processed_data['data']['policy']['internetFirewall']['policy']['sections'] = processed_sections
             
-        # Write the processed data to the new file
-        with open(output_file, 'w', encoding='utf-8') as f:
-            json.dump(processed_data, f, indent=4, ensure_ascii=False)    
-        
+        output_file = writeDataToFile(
+            data=processed_data,
+            args=args,
+            account_id=account_id,
+            default_filename_template="all_ifw_rules_and_sections_{account_id}.json",
+            default_directory="config_data"
+        )
+
         return [{"success": True, "output_file": output_file, "account_id": account_id}]
         
     except Exception as e:
@@ -206,40 +126,8 @@ def export_wf_rules_to_json(args, configuration):
     Adapted from scripts/export_wf_rules_to_json.py
     """
     try:
-        # Get account ID from args, configuration, or environment variable
-        account_id = None
-        if hasattr(args, 'accountID') and args.accountID:
-            account_id = args.accountID
-        elif hasattr(configuration, 'accountID') and configuration.accountID:
-            account_id = configuration.accountID
-        else:
-            account_id = os.getenv('CATO_ACCOUNT_ID')
-        
-        if not account_id:
-            raise ValueError("Account ID is required. Provide it using the -accountID flag or set CATO_ACCOUNT_ID environment variable.")
+        account_id = getAccountID(args, configuration)
         
-        # Set up output file path
-        if hasattr(args, 'output_file_path') and args.output_file_path:
-            # Use output file path if provided
-            output_file = args.output_file_path
-            destination_dir = os.path.dirname(output_file)
-            if hasattr(args, 'verbose') and args.verbose:
-                print(f"Using output file path: {output_file}")
-        else:
-            # Use default path and filename
-            destination_dir = 'config_data'
-            json_output_file = f"all_wf_rules_and_sections_{account_id}.json"
-            output_file = os.path.join(destination_dir, json_output_file)
-            if hasattr(args, 'verbose') and args.verbose:
-                print(f"Using default path: {output_file}")
-        
-        # Create destination directory if it doesn't exist
-        if destination_dir and not os.path.exists(destination_dir):
-            if hasattr(args, 'verbose') and args.verbose:
-                print(f"Creating directory: {destination_dir}")
-            os.makedirs(destination_dir)
-        
-        # Define the GraphQL query for WAN Firewall
         policy_query = {
             "query": "query policy ( $accountId:ID! ) { policy ( accountId:$accountId ) { wanFirewall { policy { enabled rules { audit { updatedTime updatedBy } rule { id name description index section { id name } enabled source { host { id name } site { id name } subnet ip ipRange { from to } globalIpRange { id name } networkInterface { id name } siteNetworkSubnet { id name } floatingSubnet { id name } user { id name } usersGroup { id name } group { id name } systemGroup { id name } } connectionOrigin country { id name } device { id name } deviceOS deviceAttributes { category type model manufacturer os osVersion } destination { host { id name } site { id name } subnet ip ipRange { from to } globalIpRange { id name } networkInterface { id name } siteNetworkSubnet { id name } floatingSubnet { id name } user { id name } usersGroup { id name } group { id name } systemGroup { id name } } application { application { id name } appCategory { id name } customApp { id name } customCategory { id name } sanctionedAppsCategory { id name } domain fqdn ip subnet ipRange { from to } globalIpRange { id name } } service { standard { id name } custom { port portRange { from to } protocol } } action tracking { event { enabled } alert { enabled frequency subscriptionGroup { id name } webhook { id name } mailingList { id name } } } schedule { activeOn customTimeframePolicySchedule: customTimeframe { from to } customRecurringPolicySchedule: customRecurring { from to days } } direction exceptions { name source { host { id name } site { id name } subnet ip ipRange { from to } globalIpRange { id name } networkInterface { id name } siteNetworkSubnet { id name } floatingSubnet { id name } user { id name } usersGroup { id name } group { id name } systemGroup { id name } } deviceOS destination { host { id name } site { id name } subnet ip ipRange { from to } globalIpRange { id name } networkInterface { id name } siteNetworkSubnet { id name } floatingSubnet { id name } user { id name } usersGroup { id name } group { id name } systemGroup { id name } } country { id name } device { id name } deviceAttributes { category type model manufacturer os osVersion } application { application { id name } appCategory { id name } customApp { id name } customCategory { id name } sanctionedAppsCategory { id name } domain fqdn ip subnet ipRange { from to } globalIpRange { id name } } service { standard { id name } custom { port portRangeCustomService: portRange { from to } protocol } } connectionOrigin direction } } properties } sections { audit { updatedTime updatedBy } section { id name } properties } audit { publishedTime publishedBy } revision { id name description changes createdTime updatedTime } } } } }",
             "variables": {
@@ -247,51 +135,11 @@ def export_wf_rules_to_json(args, configuration):
             },
             "operationName": "policy"
         }
-        
-        # Create API client instance with params
-        instance = CallApi(ApiClient(configuration))
-        params = {
-            'v': hasattr(args, 'verbose') and args.verbose,  # verbose mode
-            'f': 'json',  # format
-            'p': False,  # pretty print
-            't': False   # test mode
-        }
-        
-        try:
-            # Call the API directly
-            response = instance.call_api(policy_query, params)
-            all_wf_rules = response[0] if response else {}
-            
-            # Show raw API response in verbose mode
-            if hasattr(args, 'verbose') and args.verbose:
-                print("\n" + "=" * 80)
-                print("RAW API RESPONSE:")
-                print("=" * 80)
-                print(json.dumps(all_wf_rules, indent=2))
-                print("=" * 80 + "\n")
-            
-            # Check for GraphQL errors first
-            if 'errors' in all_wf_rules:
-                error_messages = [error.get('message', 'Unknown error') for error in all_wf_rules['errors']]
-                raise Exception(f"API returned errors: {', '.join(error_messages)}")
-            
-            if not all_wf_rules or 'data' not in all_wf_rules:
-                raise ValueError("Failed to retrieve data from API")
+        all_wf_rules = makeCall(args, configuration, policy_query)
                 
-        except ApiException as e:
-            raise Exception(f"API call failed - {e}")
-        except Exception as e:
-            raise Exception(f"Unexpected error during API call - {e}")
+        if not all_wf_rules or 'data' not in all_wf_rules:
+            raise ValueError("Failed to retrieve data from API")
    
-        # IMPORTANT: Preserve section IDs BEFORE stripping them
-        section_id_map = {}
-        section_to_start_after_id = None
-        sections_with_ids = all_wf_rules['data']['policy']['wanFirewall']['policy']['sections']
-        for index, section_data in enumerate(sections_with_ids):
-            section_name = section_data['section']['name']
-            section_id = section_data['section']['id']
-            section_id_map[section_name] = section_id
-            
         # Processing data to strip id attributes
         processed_data = strip_ids_recursive(all_wf_rules)
         
@@ -335,24 +183,25 @@ def export_wf_rules_to_json(args, configuration):
         
         # Reformat sections array to have index, section_id and section_name structure
         # Exclude the first section from export
+        sections_with_ids = all_wf_rules['data']['policy']['wanFirewall']['policy']['sections']
         processed_sections = [] 
-        for index, section_data in enumerate(processed_data['data']['policy']['wanFirewall']['policy']['sections']):
+        for index, section_data in enumerate(sections_with_ids):
             processed_sections.append({
                 "section_index": index,
-                "section_name": section_data['section']['name']
+                "section_name": section_data['section']['name'],
+                "section_id": section_data['section']['id']
             })
-
-        # Add preserved section IDs and section_to_start_after_id
-        processed_data['data']['policy']['wanFirewall']['policy']['section_ids'] = section_id_map
-        if section_to_start_after_id:
-            processed_data['data']['policy']['wanFirewall']['policy']['section_to_start_after_id'] = section_to_start_after_id
         
         # Replace the original sections array with the reformatted one
         processed_data['data']['policy']['wanFirewall']['policy']['sections'] = processed_sections
             
-        # Write the processed data to the new file
-        with open(output_file, 'w', encoding='utf-8') as f:
-            json.dump(processed_data, f, indent=4, ensure_ascii=False)    
+        output_file = writeDataToFile(
+            data=processed_data,
+            args=args,
+            account_id=account_id,
+            default_filename_template="all_wf_rules_and_sections_{account_id}.json",
+            default_directory="config_data"
+        )
         
         return [{"success": True, "output_file": output_file, "account_id": account_id}]
         

catocli/parsers/custom/export_sites/__init__.py

@@ -0,0 +1,20 @@
+import catocli.parsers.custom.export_sites.export_sites as export_sites
+
+def export_sites_parse(subparsers):
+    """Create export_sites command parsers"""
+    
+    # Create the socket_sites parser (direct command, no subparsers)
+    socket_sites_parser = subparsers.add_parser(
+        'socket_sites', 
+        help='Export consolidated site and socket data to JSON format',
+        usage='catocli export socket_sites [-accountID <account_id>] [options]'
+    )
+    
+    socket_sites_parser.add_argument('-accountID', help='Account ID to export data from (uses CATO_ACCOUNT_ID environment variable if not specified)', required=False)
+    socket_sites_parser.add_argument('-siteIDs', help='Comma-separated list of site IDs to export (e.g., "132606,132964,133511")', required=False)
+    socket_sites_parser.add_argument('--output-file-path', help='Full path including filename and extension for output file. If not specified, uses default: config_data/socket_site_data_{account_id}.json')
+    socket_sites_parser.add_argument('-v', '--verbose', action='store_true', help='Verbose output')
+    
+    socket_sites_parser.set_defaults(func=export_sites.export_socket_site_to_json)
+    
+    return socket_sites_parser

catocli/parsers/custom/export_sites/export_sites.py

@@ -0,0 +1,233 @@
+import os
+import json
+from graphql_client.api.call_api import ApiClient, CallApi
+from graphql_client.api_client import ApiException
+from ..customLib import writeDataToFile, makeCall, getAccountID
+
+def export_socket_site_to_json(args, configuration):
+    """
+    Export consolidated site and socket data to JSON format
+    """
+    processed_data = {'sites':{}}
+    
+    try:
+        account_id = getAccountID(args, configuration)
+        # Get account snapshot with siteIDs if provided
+        # Get siteIDs from args if provided (comma-separated string)
+        site_ids = []
+        if hasattr(args, 'siteIDs') and args.siteIDs:
+            # Parse comma-separated string into list, removing whitespace
+            site_ids = [site_id.strip() for site_id in args.siteIDs.split(',') if site_id.strip()]
+            if hasattr(args, 'verbose') and args.verbose:
+                print(f"Filtering snapshot for site IDs: {site_ids}")
+        
+        ###############################################################
+        ## Call APIs to retrieve sites, interface and network ranges ##
+        ###############################################################
+        snapshot_sites = getAccountSnapshot(args, configuration, account_id, site_ids)
+        entity_network_interfaces = getEntityLookup(args, configuration, account_id, "networkInterface")
+        entity_network_ranges = getEntityLookup(args, configuration, account_id, "siteRange")
+        entity_sites = getEntityLookup(args, configuration, account_id, "site")
+        
+        ##################################################################
+        ## Create processed_data object indexed by siteId with location ##
+        ##################################################################
+        for site_data in snapshot_sites:
+            cur_site = {
+                'wan_interfaces': {},
+                'lan_interfaces': {},
+            }
+            site_id = site_data.get('id')
+            cur_site['id'] = site_id
+            cur_site['name'] = site_data.get('infoSiteSnapshot', {}).get('name')
+            cur_site['description'] = site_data.get('infoSiteSnapshot', {}).get('description')
+            cur_site['connectionType'] = site_data.get('infoSiteSnapshot', {}).get('connType')
+            cur_site['type'] = site_data.get('infoSiteSnapshot', {}).get('type')
+            cur_site = populateSiteLocationData(args, site_data, cur_site)
+
+            site_interfaces = site_data.get('infoSiteSnapshot', {}).get('interfaces', [])
+            for wan_ni in site_interfaces:
+                cur_wan_interface = {}
+                id = wan_ni.get('wanRoleInterfaceInfo', "")
+                if id[0:3] == "wan":
+                    cur_wan_interface['id'] = id
+                    cur_wan_interface['name'] = wan_ni.get('name', "")
+                    cur_wan_interface['upstreamBandwidth'] = wan_ni.get('upstreamBandwidth', 0)
+                    cur_wan_interface['downstreamBandwidth'] = wan_ni.get('downstreamBandwidth', 0)
+                    cur_wan_interface['destType'] = wan_ni.get('destType', "")
+                    cur_site['wan_interfaces'][id] = cur_wan_interface
+
+            if site_id:
+                processed_data['sites'][site_id] = cur_site
+
+        ##################################################################################
+        ## Process entity lookup LAN network interfaces adding to site object by site_id##
+        ##################################################################################
+        for lan_ni in entity_network_interfaces:
+            cur_lan_interface = {
+                'network_ranges': {},
+            }
+            site_id = lan_ni.get("helperFields","").get('siteId', "")
+            id = lan_ni.get('entity', "").get('id', "")
+            interfaceName = lan_ni.get('entity', "").get('interfaceName', "")
+            cur_lan_interface['id'] = id
+            cur_lan_interface['name'] = interfaceName
+            cur_lan_interface['index'] = lan_ni.get("helperFields","").get('interfaceId', "")
+            cur_lan_interface['upstreamBandwidth'] = lan_ni.get('upstreamBandwidth', 0)
+            cur_lan_interface['downstreamBandwidth'] = lan_ni.get('downstreamBandwidth', 0)
+            cur_lan_interface['destType'] = lan_ni.get('destType', "")
+            processed_data['sites'][site_id]['lan_interfaces'][interfaceName] = cur_lan_interface
+
+        #############################################################################
+        ## Process entity lookup network ranges populating by network interface id ##
+        #############################################################################
+        for range in entity_network_ranges:
+            cur_range = {}
+            site_id = lan_ni.get("helperFields","").get('siteId', "")
+            id = lan_ni.get('entity', "").get('id', "")
+            interface_name = lan_ni.get('entity', "").get('interfaceName', "")
+            cur_lan_interface['id'] = id
+            cur_lan_interface['subnet'] = lan_ni.get("helperFields","").get('subnet', "")
+            cur_lan_interface['vlanTag'] = lan_ni.get("helperFields","").get('vlanTag', "")
+            cur_lan_interface['microsegmentation'] = lan_ni.get("helperFields","").get('microsegmentation', "")
+            
+            processed_data['sites'][site_id]['lan_interfaces'][interface_name] = cur_range
+
+
+        # Write the processed data to file using the general-purpose function
+        output_file = writeDataToFile(
+            data=processed_data,
+            args=args,
+            account_id=account_id,
+            default_filename_template="socket_sites_{account_id}.json",
+            default_directory="config_data"
+        )
+        
+        return [{"success": True, "output_file": output_file, "account_id": account_id}]
+            
+    except Exception as e:
+        print(f"ERROR: {str(e)}")
+        return [{"success": False, "error": str(e)}]
+
+
+##########################################################################
+########################### Helper functions #############################
+##########################################################################
+
+def populateSiteLocationData(args, site_data, cur_site):
+    # Load site location data for timezone and state code lookups
+    site_location_data = {}
+    try:
+        script_dir = os.path.dirname(os.path.abspath(__file__))
+        models_dir = os.path.join(script_dir, '..', '..', '..', '..', 'models')
+        location_file = os.path.join(models_dir, 'query.siteLocation.json')
+        
+        if os.path.exists(location_file):
+            with open(location_file, 'r', encoding='utf-8') as f:
+                site_location_data = json.load(f)
+            if hasattr(args, 'verbose') and args.verbose:
+                print(f"Loaded {len(site_location_data)} location entries from {location_file}")
+        else:
+            if hasattr(args, 'verbose') and args.verbose:
+                print(f"Warning: Site location file not found at {location_file}")
+    except Exception as e:
+        if hasattr(args, 'verbose') and args.verbose:
+            print(f"Warning: Could not load site location data: {e}")
+
+    ## siteLocation attributes
+    cur_site['address'] = site_data.get('infoSiteSnapshot', {}).get('address')
+    cur_site['city'] = site_data.get('infoSiteSnapshot', {}).get('cityName')                
+    cur_site['stateName'] = site_data.get('infoSiteSnapshot', {}).get('countryStateName')
+    cur_site['countryCode'] = site_data.get('infoSiteSnapshot', {}).get('countryCode')
+    cur_site['countryName'] = site_data.get('infoSiteSnapshot', {}).get('countryName')
+
+    # Look up timezone and state code from location data
+    country_name = cur_site['countryName']
+    state_name = cur_site['stateName']
+    city = cur_site['city']
+
+    # Create lookup key based on available data
+    if state_name:
+        lookup_key = f"{country_name}___{state_name}___{city}"
+    else:
+        lookup_key = f"{country_name}___{city}"
+    
+    # Debug output for lookup
+    if hasattr(args, 'verbose') and args.verbose:
+        print(f"Site {cur_site['name']}: Looking up '{lookup_key}'")
+
+    # Look up location details
+    location_data = site_location_data.get(lookup_key, {})
+    
+    if hasattr(args, 'verbose') and args.verbose:
+        if location_data:
+            print(f"  Found location data: {location_data}")
+        else:
+            print(f"  No location data found for key: {lookup_key}")
+            # Try to find similar keys for debugging
+            similar_keys = [k for k in site_location_data.keys() if country_name in k and (not city or city in k)][:5]
+            if similar_keys:
+                print(f"  Similar keys found: {similar_keys}")
+
+    cur_site['stateCode'] = location_data.get('stateCode', None)
+
+    # Get timezone - always use the 0 element in the timezones array
+    timezones = location_data.get('timezone', [])
+    cur_site['timezone'] = timezones[0] if timezones else None
+    return cur_site
+
+def getEntityLookup(args, configuration, account_id, entity_type):
+    """
+    Helper function to get entity lookup data for a specific entity type
+    """
+    #################################
+    ## Get entity lookup for sites ##
+    #################################
+    entity_query = {
+        "query": "query entityLookup ( $accountID:ID! $type:EntityType! $sortInput:[SortInput] $lookupFilterInput:[LookupFilterInput] ) { entityLookup ( accountID:$accountID type:$type sort:$sortInput filters:$lookupFilterInput ) { items { entity { id name type } description helperFields } total } }",
+        "variables": {
+            "accountID": account_id,
+            "type": entity_type
+        },
+        "operationName": "entityLookup"
+    }
+    response = makeCall(args, configuration, entity_query)
+
+    # Check for GraphQL errors in snapshot response
+    if 'errors' in response:
+        error_messages = [error.get('message', 'Unknown error') for error in response['errors']]
+        raise Exception(f"Snapshot API returned errors: {', '.join(error_messages)}")
+    
+    if not response or 'data' not in response or 'entityLookup' not in response['data']:
+        raise ValueError("Failed to retrieve snapshot data from API")
+    
+    items = response['data']['entityLookup']['items']
+    if items is None:
+        items = []
+        if hasattr(args, 'verbose') and args.verbose:
+            print("No items found in entity lookup - "+ entity_type)
+    return items
+
+def getAccountSnapshot(args, configuration, account_id, site_ids=None):
+    snapshot_query = {
+        "query": "query accountSnapshot ( $siteIDs:[ID!] $accountID:ID ) { accountSnapshot ( accountID:$accountID ) { id sites ( siteIDs:$siteIDs ) { id protoId connectivityStatusSiteSnapshot: connectivityStatus haStatusSiteSnapshot: haStatus { readiness wanConnectivity keepalive socketVersion } operationalStatusSiteSnapshot: operationalStatus lastConnected connectedSince popName devices { id name identifier connected haRole interfaces { connected id name physicalPort naturalOrder popName previousPopID previousPopName tunnelConnectionReason tunnelUptime tunnelRemoteIP tunnelRemoteIPInfoInterfaceSnapshot: tunnelRemoteIPInfo { ip countryCode countryName city state provider latitude longitude } type infoInterfaceSnapshot: info { id name upstreamBandwidth downstreamBandwidth upstreamBandwidthMbpsPrecision downstreamBandwidthMbpsPrecision destType wanRole } cellularInterfaceInfoInterfaceSnapshot: cellularInterfaceInfo { networkType simSlotId modemStatus isModemConnected iccid imei operatorName isModemSuspended apn apnSelectionMethod signalStrength isRoamingAllowed simNumber disconnectionReason isSimSlot1Detected isSimSlot2Detected } } lastConnected lastDuration connectedSince lastPopID lastPopName recentConnections { duration interfaceName deviceName lastConnected popName remoteIP remoteIPInfoRecentConnection: remoteIPInfo { ip countryCode countryName city state provider latitude longitude } } type deviceUptime socketInfo { id serial isPrimary platformSocketInfo: platform version versionUpdateTime } interfacesLinkState { id up mediaIn linkSpeed duplex hasAddress hasInternet hasTunnel } osType osVersion version versionNumber releaseGroup mfaExpirationTime mfaCreationTime internalIP } infoSiteSnapshot: info { name type description countryCode region countryName countryStateName cityName address isHA connType creationTime interfaces { id name upstreamBandwidth downstreamBandwidth upstreamBandwidthMbpsPrecision downstreamBandwidthMbpsPrecision destType wanRoleInterfaceInfo: wanRole } sockets { id serial isPrimary platformSocketInfo: platform version versionUpdateTime } ipsec { isPrimary catoIP remoteIP ikeVersion } } hostCount altWanStatus } users { id connectivityStatusUserSnapshot: connectivityStatus operationalStatusUserSnapshot: operationalStatus name deviceName uptime lastConnected version versionNumber popID popName remoteIP remoteIPInfoUserSnapshot: remoteIPInfo { ip countryCode countryName city state provider latitude longitude } internalIP osType osVersion devices { id name identifier connected haRole interfaces { connected id name physicalPort naturalOrder popName previousPopID previousPopName tunnelConnectionReason tunnelUptime tunnelRemoteIP tunnelRemoteIPInfoInterfaceSnapshot: tunnelRemoteIPInfo { ip countryCode countryName city state provider latitude longitude } type infoInterfaceSnapshot: info { id name upstreamBandwidth downstreamBandwidth upstreamBandwidthMbpsPrecision downstreamBandwidthMbpsPrecision destType wanRole } cellularInterfaceInfoInterfaceSnapshot: cellularInterfaceInfo { networkType simSlotId modemStatus isModemConnected iccid imei operatorName isModemSuspended apn apnSelectionMethod signalStrength isRoamingAllowed simNumber disconnectionReason isSimSlot1Detected isSimSlot2Detected } } lastConnected lastDuration connectedSince lastPopID lastPopName recentConnections { duration interfaceName deviceName lastConnected popName remoteIP remoteIPInfoRecentConnection: remoteIPInfo { ip countryCode countryName city state provider latitude longitude } } type deviceUptime socketInfo { id serial isPrimary platformSocketInfo: platform version versionUpdateTime } interfacesLinkState { id up mediaIn linkSpeed duplex hasAddress hasInternet hasTunnel } osType osVersion version versionNumber releaseGroup mfaExpirationTime mfaCreationTime internalIP } connectedInOffice infoUserSnapshot: info { name status email creationTime phoneNumber origin authMethod } recentConnections { duration interfaceName deviceName lastConnected popName remoteIP remoteIPInfo { ip countryCode countryName city state provider latitude longitude } } } timestamp }  }",
+        "variables": {
+            "accountID": account_id,
+            "siteIDs": site_ids
+        },
+        "operationName": "accountSnapshot"
+    }
+    response = makeCall(args, configuration, snapshot_query)
+
+    # Check for GraphQL errors in snapshot response
+    if 'errors' in response:
+        error_messages = [error.get('message', 'Unknown error') for error in response['errors']]
+        raise Exception(f"Snapshot API returned errors: {', '.join(error_messages)}")
+    
+    if not response or 'data' not in response or 'accountSnapshot' not in response['data']:
+        raise ValueError("Failed to retrieve snapshot data from API")
+    
+    if not response or 'sites' not in response['data']['accountSnapshot'] or response['data']['accountSnapshot']['sites'] is None:
+        raise ValueError("No sites found in account snapshot data from API")
+
+    return response

catocli/parsers/custom/import_rules_to_tf/__init__.py

@@ -11,7 +11,7 @@ def import_parse(subparsers):
     if_rules_parser = import_subparsers.add_parser(
         'if_rules_to_tf', 
         help='Import Internet Firewall rules to Terraform state',
-        usage='catocli import if_rules_to_tf <json_file> --module-name <module_name> [options]'
+        usage='catocli import if_rules_to_tf <json_file> --module-name <module_name> [options]\n\nexample: catocli import if_rules_to_tf config_data/all_wf_rules_and_sections.json --module-name module.if_rules'
     )
     
     if_rules_parser.add_argument('json_file', help='Path to the JSON file containing IFW rules and sections')
@@ -35,7 +35,7 @@ def import_parse(subparsers):
     wf_rules_parser = import_subparsers.add_parser(
         'wf_rules_to_tf', 
         help='Import WAN Firewall rules to Terraform state',
-        usage='catocli import wf_rules_to_tf <json_file> --module-name <module_name> [options]'
+        usage='catocli import wf_rules_to_tf <json_file> --module-name <module_name> [options]\n\nexample: catocli import wf_rules_to_tf config_data/all_wf_rules_and_sections.json --module-name module.wf_rules'
     )
     
     wf_rules_parser.add_argument('json_file', help='Path to the JSON file containing WF rules and sections')

catocli/parsers/custom/import_rules_to_tf/import_rules_to_tf.py

@@ -60,14 +60,13 @@ def extract_rules_and_sections(policy_data):
             })
     
     # Extract sections
-    section_ids = policy_data.get('section_ids', {})
     for section in policy_data.get('sections', []):
         if section.get('section_name'):
             sections.append({
                 'section_name': section['section_name'],
                 'section_index': section.get('section_index', 0),
-                'section_id': section_ids.get(section['section_name'], '')
-            })  
+                'section_id': section.get('section_id', '')
+            })
     return rules, sections
 
 
@@ -123,8 +122,8 @@ def find_rule_index(rules, rule_name):
     return None
 
 
-def import_sections(sections, module_name, verbose=False,
-                   resource_type="cato_if_section", resource_name="sections"):
+def import_sections(sections, module_name, resource_type,
+                    resource_name="sections", verbose=False):
     """Import all sections"""
     print("\nStarting section imports...")
     total_sections = len(sections)
@@ -135,7 +134,7 @@ def import_sections(sections, module_name, verbose=False,
         section_id = section['section_id']
         section_name = section['section_name']
         section_index = section['section_index']
-        resource_address = f'{module_name}.{resource_type}.{resource_name}["{str(section_name)}"]'
+        resource_address = f'{module_name}.{resource_type}.{resource_name}["{section_name}"]'
         print(f"\n[{i+1}/{total_sections}] Section: {section_name} (index: {section_index})")
 
         # For sections, we use the section name as the ID since that's how Cato identifies them
@@ -376,7 +375,7 @@ def import_if_rules_to_tf(args, configuration):
         
         # Import sections first (if not skipped)
         if not args.rules_only and sections:
-            successful, failed = import_sections(sections, module_name=args.module_name, verbose=args.verbose)
+            successful, failed = import_sections(sections, module_name=args.module_name, resource_type="cato_if_section", verbose=args.verbose)
             total_successful += successful
             total_failed += failed
         
@@ -439,7 +438,10 @@ def import_wf_sections(sections, module_name, verbose=False,
         section_id = section['section_id']
         section_name = section['section_name']
         section_index = section['section_index']
-        resource_address = f'{module_name}.{resource_type}.{resource_name}["{str(section_name)}"]'
+        # Add module. prefix if not present
+        if not module_name.startswith('module.'):
+            module_name = f'module.{module_name}'
+        resource_address = f'{module_name}.{resource_type}.{resource_name}["{section_name}"]'
         print(f"\n[{i+1}/{total_sections}] Section: {section_name} (index: {section_index})")
 
         # For sections, we use the section name as the ID since that's how Cato identifies them
@@ -469,6 +471,10 @@ def import_wf_rules(rules, module_name, verbose=False,
         rule_index = find_rule_index(rules, rule_name)
         terraform_key = sanitize_name_for_terraform(rule_name)
         
+        # Add module. prefix if not present
+        if not module_name.startswith('module.'):
+            module_name = f'module.{module_name}'
+
         # Use array index syntax instead of rule ID
         resource_address = f'{module_name}.{resource_type}.{resource_name}["{str(rule_name)}"]'
         print(f"\n[{i+1}/{total_rules}] Rule: {rule_name} (index: {rule_index})")
@@ -519,8 +525,12 @@ def import_wf_rules_to_tf(args, configuration):
             print(" No rules or sections found. Exiting.")
             return [{"success": False, "error": "No rules or sections found"}]
         
+        # Add module. prefix if not present
+        module_name = args.module_name
+        if not module_name.startswith('module.'):
+            module_name = f'module.{module_name}'
         # Validate Terraform environment before proceeding
-        validate_terraform_environment(args.module_name, verbose=args.verbose)
+        validate_terraform_environment(module_name, verbose=args.verbose)
         
         # Ask for confirmation (unless auto-approved)
         if not args.rules_only and not args.sections_only: