regscale-cli 6.21.0.0__py3-none-any.whl → 6.21.2.0__py3-none-any.whl

regscale/integrations/commercial/qualys/__init__.py

@@ -36,6 +36,7 @@ from regscale.integrations.scanner_integration import IntegrationAsset, Integrat
 from regscale.integrations.variables import ScannerVariables
 from regscale.models import Asset, Issue, Search, regscale_models
 from regscale.models.app_models.click import NotRequiredIf, regscale_ssp_id, save_output_to, ssp_or_component_id
+from regscale.models.integration_models.ecr_models.data import Finding
 from regscale.models.integration_models.flat_file_importer import FlatFileImporter
 from regscale.models.integration_models.qualys import (
     Qualys,
@@ -73,9 +74,9 @@ def _prepare_qualys_params(include_tags: str, exclude_tags: str) -> dict:
         "show_asset_id": "1",
         "show_tags": "1",
     }
-
     if exclude_tags or include_tags:
         params["use_tags"] = "1"
+        params["tag_set_by"] = "name"
         if exclude_tags:
             params["tag_set_exclude"] = exclude_tags
         if include_tags:
@@ -179,12 +180,12 @@ def _log_asset_results(count, asset_ids):
     :param asset_ids: List of asset IDs
     """
     if count == 0:
-        logger.warning("No assets were yielded from the JSONL file")
+        logger.warning("No assets were created/updated from the JSONL file")
     else:
         sample_ids = asset_ids[:5]
         truncation_indicator = ", ..." if len(asset_ids) > 5 else ""
         sample_ids_str = ", ".join(sample_ids)
-        logger.debug(f"Yielded {count} assets to RegScale with IDs: {sample_ids_str}{truncation_indicator}")
+        logger.debug(f"Created/updated {count} assets to RegScale with IDs: {sample_ids_str}{truncation_indicator}")
 
 
 def _import_findings(integration, progress, finding_task):
@@ -270,6 +271,7 @@ class FindingProgressTracker:
         self.finding_task = finding_task
         self.count = 0
         self.finding_ids = []
+        self.output_final_log: bool = False
 
     def __iter__(self):
         return self
@@ -284,22 +286,22 @@ class FindingProgressTracker:
             return finding
         except StopIteration:
             self._log_finding_results()
-            raise
+            return
         except Exception as e:
-            logger.error(f"Error yielding findings: {str(e)}")
-            logger.debug(f"Findings yielded before error: {self.count}")
-            raise
+            logger.debug(f"Findings created/updated before error: {self.count}")
+            error_and_exit(f"Error creating/updating findings: {str(e)}")
 
     def _log_finding_results(self):
         """Log the results of finding processing."""
         if self.count == 0:
-            logger.warning("No findings were yielded from the JSONL file")
-        else:
-            logger.info(f"Yielded {self.count} findings to RegScale")
+            logger.warning("No findings were created/updated from the JSONL file")
+        elif not self.output_final_log:
+            logger.info(f"Created/Updated {self.count} findings to RegScale")
             sample_ids = self.finding_ids[:5]
             truncation_indicator = ", ..." if len(self.finding_ids) > 5 else ""
             sample_ids_str = ", ".join(sample_ids)
             logger.debug(f"Sample finding IDs: {sample_ids_str}{truncation_indicator}")
+            self.output_final_log = True
 
 
 @click.command(name="import_total_cloud")
@@ -318,7 +320,7 @@ class FindingProgressTracker:
     type=click.STRING,
     required=False,
     default=None,
-    help="Exclude tags in the import comma separated string of tag names or ids, defaults to None.",
+    help="Exclude tags in the import comma separated string of tag names or ids, defaults to None. If used, --include_tags must also be provided.",
 )
 @click.option(
     "--vulnerability-creation",
@@ -368,10 +370,14 @@ def import_total_cloud(
             logger.warning("SSP #%i is not a valid RegScale Security Plan.", regscale_ssp_id)
             return
     else:
-        raise click.UsageError(
+        error_and_exit(
             "You must provide either a --regscale_ssp_id or a --component_id to import Qualys Total Cloud data."
         )
 
+    # exclude tags must have include_tags
+    if exclude_tags and not include_tags:
+        error_and_exit("You must provide --include_tags when using --exclude_tags to import Qualys Total Cloud data.")
+
     containers_lst = []
     try:
         # Configure scanner variables and fetch data
@@ -382,7 +388,8 @@ def import_total_cloud(
 
         if containers:
             # Fetch containers and container findings
-            containers_lst = fetch_all_vulnerabilities()
+            params = _prepare_qualys_params(include_tags, exclude_tags)
+            containers_lst = fetch_all_vulnerabilities(filters=params)
 
         # Initialize and run integration
         integration = _initialize_integration(
@@ -599,7 +606,7 @@ def import_total_cloud_from_xml(xml_file: str, regscale_ssp_id: int = None, comp
                 logger.warning("SSP #%i is not a valid RegScale Security Plan.", regscale_ssp_id)
                 return
         else:
-            raise click.UsageError(
+            error_and_exit(
                 "You must provide either a --regscale_ssp_id or a --component_id to import Qualys Total Cloud data."
             )
 
@@ -734,10 +741,10 @@ def _track_progress_generator(iterator, progress_bar, task_id, id_attribute=None
     elif id_attribute == "external_id":
         item_type = "findings"
 
-    logger.debug(f"Yielded {processed} {item_type} to RegScale")
+    logger.debug(f"Created/updated {processed} {item_type} in RegScale")
 
     if processed == 0:
-        logger.error(f"WARNING: No {item_type} were yielded to RegScale!")
+        logger.error(f"WARNING: No {item_type} were created/updated in RegScale!")
     elif item_ids:
         logger.debug(f"First 10 {item_type} IDs: {item_ids[:10]}")
 
@@ -978,7 +985,7 @@ def _resolve_plan_and_component(regscale_ssp_id: int = None, component_id: int =
     Returns (plan_id, is_component)
     """
     if (regscale_ssp_id is None and component_id is None) or (regscale_ssp_id and component_id):
-        raise click.UsageError("You must provide either --regscale_ssp_id or --component_id, but not both.")
+        error_and_exit("You must provide either --regscale_ssp_id or --component_id, but not both.")
     is_component = component_id is not None
     plan_id = component_id if is_component else regscale_ssp_id
     return plan_id, is_component
@@ -1455,75 +1462,6 @@ def _get_qualys_api():
     return qualys_url, QUALYS_API
 
 
-def import_total_cloud_data_from_qualys_api(security_plan_id: int, include_tags: str, exclude_tags: str):
-    """
-    Function to get the total cloud data from Qualys API
-
-    :param int security_plan_id: The ID of the plan to get the data for
-    :param str include_tags: The tags to include in the data
-    :param str exclude_tags: The tags to exclude from the data
-    """
-    try:
-        qualys_url, QUALYS_API = _get_qualys_api()
-        params = {
-            "action": "list",
-            "show_asset_id": "1",
-            "show_tags": "1",
-        }
-        if exclude_tags or include_tags:
-            params["use_tags"] = "1"
-            if exclude_tags:
-                params["tag_set_exclude"] = exclude_tags
-            if include_tags:
-                params["tag_set_include"] = include_tags
-
-        logger.info("Fetching Qualys Total Cloud data...")
-        response = QUALYS_API.get(
-            url=urljoin(qualys_url, "/api/2.0/fo/asset/host/vm/detection/"),
-            headers=HEADERS,
-            params=params,
-        )
-
-        if response and response.ok:
-            logger.info("Total cloud data fetched. processing...")
-            response_data = xmltodict.parse(response.text)
-
-            # Create artifacts directory if it doesn't exist
-            os.makedirs("./artifacts", exist_ok=True)
-
-            # Write the XML data to a temporary file
-            temp_xml_file = "./artifacts/qualys_temp_data.xml"
-            with open(temp_xml_file, "w") as f:
-                f.write(response.text)
-
-            logger.info(f"Saved Qualys response data to {temp_xml_file}")
-
-            # Initialize the JSONLScannerIntegration implementation with the file path
-            integration = QualysTotalCloudJSONLIntegration(
-                plan_id=security_plan_id, xml_data=response_data, file_path=temp_xml_file
-            )
-
-            # Process assets and findings in one pass for efficiency
-            assets_iterator, findings_iterator = integration.fetch_assets_and_findings()
-
-            # Update RegScale with the processed data
-            logger.info("Syncing assets to RegScale...")
-            integration.update_regscale_assets(assets_iterator)
-
-            logger.info("Syncing findings to RegScale...")
-            integration.update_regscale_findings(findings_iterator)
-
-            logger.info("Qualys Total Cloud data imported successfully.")
-        else:
-            logger.error(
-                f"Received unexpected response from Qualys API: {response.status_code}: {response.text if response.text else 'response is null'}"
-            )
-    except Exception:
-        error_message = traceback.format_exc()
-        logger.error("Error occurred while processing Qualys data")
-        logger.error(error_message)
-
-
 def get_scans_summary(scan_choice: str) -> dict:
     """
     Get all scans from Qualys Host

regscale/integrations/commercial/qualys/containers.py

@@ -213,6 +213,8 @@ def _fetch_paginated_data(endpoint: str, filters: Optional[Dict] = None, limit:
 
                 # Clear params for subsequent requests since they're in the URL
                 params = {}
+            progress.update(task, total=len(all_items))
+            progress.update(task, completed=len(all_items))
 
     except Exception as e:
         logger.error("Error fetching data from %s: %s", current_url if current_url else "N/A", e)

regscale/integrations/commercial/qualys/scanner.py

@@ -1096,6 +1096,11 @@ class QualysTotalCloudJSONLIntegration(JSONLScannerIntegration):
                 first_found = self._convert_timestamp_to_date_str(first_found)
                 last_found = self._convert_timestamp_to_date_str(last_found)
 
+                # sometimes cvss3Info is not a dict, so we ensure it is
+                cvs3_info = vuln.get("cvss3Info")
+                if not isinstance(cvs3_info, dict):
+                    cvs3_info = {}
+
                 cve = next(iter(vuln.get("cveids", [])), "")
                 # Create finding object
                 return IntegrationFinding(
@@ -1109,8 +1114,8 @@ class QualysTotalCloudJSONLIntegration(JSONLScannerIntegration):
                     category="Vulnerability",
                     plugin_name=cve or f"QID-{qid}",
                     control_labels=[f"QID-{qid}"],
-                    cvss_v3_base_score=vuln.get("cvss3Info", {}).get("baseScore"),
-                    cvss_v3_vector=vuln.get("cvss3Info", {}).get("temporalScore"),
+                    cvss_v3_base_score=cvs3_info.get("baseScore"),
+                    cvss_v3_vector=cvs3_info.get("temporalScore"),
                     first_seen=first_found,
                     last_seen=last_found,
                     evidence=vuln.get("result", "No evidence available"),

regscale/integrations/commercial/sonarcloud.py

@@ -13,6 +13,7 @@ import requests  # type: ignore
 from regscale.core.app.api import Api
 from regscale.core.app.application import Application
 from regscale.core.app.utils.app_utils import (
+    create_progress_object,
     days_between,
     error_and_exit,
     get_current_datetime,
@@ -24,12 +25,16 @@ from regscale.models.regscale_models.issue import Issue
 logger = logging.getLogger("regscale")
 
 
-def get_sonarcloud_results(config: dict, branch: Optional[str] = None) -> list[list[dict]]:
+def get_sonarcloud_results(
+    config: dict, organization: Optional[str] = None, branch: Optional[str] = None, project_key: Optional[str] = None
+) -> list[list[dict]]:
     """
     Retrieve Sonarcloud Results from the Sonarcloud.io API
 
     :param dict config: RegScale CLI configuration
+    :param Optional[str] organization: Organization name to filter results, defaults to None
     :param Optional[str] branch: Branch name to filter results, defaults to None
+    :param Optional[str] project_key: SonarCloud Project Key, defaults to None
     :return: json response data from API GET request
     :rtype: list[list[dict]]
     """
@@ -44,10 +49,19 @@ def get_sonarcloud_results(config: dict, branch: Optional[str] = None) -> list[l
         "statuses": "OPEN, CONFIRMED, REOPENED",
         "ps": 500,
     }
+    if organization and project_key:
+        params["componentKeys"] = project_key
+    if organization:
+        params["organization"] = organization
     if branch:
         params["branch"] = branch
+    if project_key:
+        params["projectKeys"] = project_key
     # GET request pulls in data to check results size
+    logger.info("Fetching issues from SonarCloud/Qube...")
     r = requests.get(url, auth=(str(token), ""), params=params)
+    if r.status_code != 200:
+        error_and_exit(f"Sonarcloud API call failed with status code {r.status_code}: {r.reason}\n{r.text}")
     # if the status code does not equal 200
     if r and not r.ok:
         # exit the script gracefully
@@ -57,13 +71,12 @@ def get_sonarcloud_results(config: dict, branch: Optional[str] = None) -> list[l
     # find the total results number
     total = data["paging"]["total"]
     complete.extend(data.get("issues", []))
-    logger.info(f"Found {total} issue(s) from SonarCloud/Qube.")
     # find the number of results in each result page
     size = data["paging"]["pageSize"]
     # calculate the number of pages to iterate through sequentially
     pages = math.ceil(total / size)
     # loop through each page number
-    for i in range(1, pages + 1, 1):
+    for i in range(2, pages + 1, 1):
         # parameters to pass to the API call
         params["p"] = str(i)
         # for each page make a GET request to pull in the data
@@ -79,56 +92,56 @@ def get_sonarcloud_results(config: dict, branch: Optional[str] = None) -> list[l
     return complete
 
 
-def build_data(api: Api, branch: Optional[str] = None) -> list[dict]:
+def build_data(
+    api: Api, organization: Optional[str] = None, branch: Optional[str] = None, project_key: Optional[str] = None
+) -> list[dict]:
     """
     Build vulnerability alert data list
+
     :param Api api: API object
+    :param Optional[str] organization: Organization name to filter results, defaults to None
     :param Optional[str] branch: Branch name to filter results, defaults to None
+    :param Optional[str] project_key: SonarCloud Project Key, defaults to None
     :return: vulnerability data list
     :rtype: list[dict]
     """
     # execute GET request
-    data = get_sonarcloud_results(config=api.config, branch=branch)
+    data = get_sonarcloud_results(config=api.config, organization=organization, branch=branch, project_key=project_key)
     # create empty list to hold json response dicts
     vulnerability_data_list = []
     # loop through the lists in API response data
-    for result in data:
+    for issue in data:
         # loop through the list of dicts in the API response data
-        for i, issue in enumerate(result):
-            # format datetime stamp to use with days_between function
-            create_date = issue["creationDate"][0:19] + "Z"
-            # build vulnerability list
-            vulnerability_data_list.append(
-                {
-                    "key": issue["key"],
-                    "severity": issue["severity"],
-                    "component": issue["component"],
-                    "status": issue["status"],
-                    "message": issue["message"],
-                    "creationDate": issue["creationDate"][0:19],
-                    "updateDate": issue["updateDate"][0:19],
-                    "type": issue["type"],
-                    "days_elapsed": days_between(vuln_time=create_date),
-                }
-            )
+        # format datetime stamp to use with days_between function
+        create_date = issue["creationDate"][0:19] + "Z"
+        # build vulnerability list
+        vulnerability_data_list.append(
+            {
+                "key": issue["key"],
+                "severity": issue["severity"],
+                "component": issue["component"],
+                "status": issue["status"],
+                "message": issue["message"],
+                "creationDate": issue["creationDate"][0:19],
+                "updateDate": issue["updateDate"][0:19],
+                "type": issue["type"],
+                "days_elapsed": days_between(vuln_time=create_date),
+            }
+        )
     return vulnerability_data_list
 
 
-def build_dataframes(api: Api) -> str:
+def build_dataframes(sonar_data: list[dict]) -> str:
     """
     Build pandas dataframes from vulnerability alert data list
 
-    :param Api api: API object
+    :param list[dict] sonar_data: SonarCloud alerts and issues data
     :return: dataframe as an HTML table
     :rtype: str
     """
     import pandas as pd  # Optimize import performance
 
-    # create vulnerability data list
-    vuln_data_list = build_data(api=api)
-
-    # for vulnerability in vuln_data_list:
-    df = pd.DataFrame(vuln_data_list)
+    df = pd.DataFrame(sonar_data)
     # sort dataframe by severity
     df.sort_values(by=["severity"], inplace=True)
     # reset and drop the index
@@ -139,10 +152,12 @@ def build_dataframes(api: Api) -> str:
 
 
 def create_alert_assessment(
-    api: Api, parent_id: Optional[int] = None, parent_module: Optional[str] = None
+    sonar_data: list[dict], api: Api, parent_id: Optional[int] = None, parent_module: Optional[str] = None
 ) -> Optional[int]:
     """
     Create Assessment containing SonarCloud alerts
+
+    :param list[dict] sonar_data: SonarCloud alerts and issues data
     :param Api api: API object
     :param Optional[int] parent_id: Parent ID of the assessment, defaults to None
     :param Optional[str] parent_module: Parent module of the assessment, defaults to None
@@ -150,7 +165,7 @@ def create_alert_assessment(
     :rtype: Optional[int]
     """
     # create the assessment report HTML table
-    df_output = build_dataframes(api)
+    df_output = build_dataframes(sonar_data)
     # build assessment model data
     assessment_data = Assessment(
         leadAssessorId=api.config["userId"],
@@ -169,11 +184,9 @@ def create_alert_assessment(
     if parent_id and parent_module:
         assessment_data.parentId = parent_id
         assessment_data.parentModule = parent_module
-    # create vulnerability data list
-    vuln_data_list = build_data(api)
     # if assessmentResult is changed to Pass / Fail then status has to be
     # changed to complete and a completion date has to be passed
-    for vulnerability in vuln_data_list:
+    for vulnerability in sonar_data:
         if vulnerability["severity"] == "CRITICAL" and vulnerability["days_elapsed"] >= 10:
             assessment_data.status = "Complete"
             assessment_data.actualFinish = get_current_datetime()
@@ -182,60 +195,68 @@ def create_alert_assessment(
     # create a new assessment in RegScale
     if new_assessment := assessment_data.create():
         # log assessment creation result
-        api.logger.info("Assessment was created successfully")
+        api.logger.debug("Assessment was created successfully")
         return new_assessment.id
     else:
-        api.logger.info("Assessment was not created")
+        api.logger.debug("Assessment was not created")
         return None
 
 
 def create_alert_issues(
-    parent_id: Optional[int] = None, parent_module: Optional[str] = None, branch: Optional[str] = None
+    parent_id: Optional[int] = None,
+    parent_module: Optional[str] = None,
+    organization: Optional[str] = None,
+    branch: Optional[str] = None,
+    project_key: Optional[str] = None,
 ) -> None:
     """
     Create child issues from the alert assessment
+
     :param Optional[int] parent_id: Parent ID record to associate the assessment to, defaults to None
     :param Optional[str] parent_module: Parent module to associate the assessment to, defaults to None
+    :param Optional[str] organization: Organization name to filter results, defaults to None
     :param Optional[str] branch: Branch name to filter results, defaults to None
+    :param Optional[str] project_key: SonarCloud Project Key, defaults to None
     :rtype: None
     """
     # set environment and application configuration
     app = Application()
     api = Api()
+    sonar_data = build_data(api=api, organization=organization, branch=branch, project_key=project_key)
     # execute POST request and return new assessment ID
-    assessment_id = create_alert_assessment(api=api, parent_id=parent_id, parent_module=parent_module)
+    assessment_id = create_alert_assessment(
+        sonar_data=sonar_data, api=api, parent_id=parent_id, parent_module=parent_module
+    )
 
     # create vulnerability data list
-    vuln_data_list = build_data(api, branch)
     # loop through each vulnerability alert in the list
-    for vulnerability in vuln_data_list:
-        # create issue model
-        issue_data = Issue(
-            title="Sonarcloud Code Scan",  # Required
-            dateCreated=get_current_datetime("%Y-%m-%dT%H:%M:%S"),
-            description=vulnerability["message"],
-            severityLevel=Issue.assign_severity(vulnerability["severity"]),  # Required
-            issueOwnerId=app.config["userId"],  # Required
-            dueDate=get_current_datetime(),
-            identification="Code scan assessment",
-            status="Open",
-            assessmentId=assessment_id,
-            createdBy=app.config["userId"],
-            lastUpdatedById=app.config["userId"],
-            dateLastUpdated=get_current_datetime(),
-            parentId=assessment_id,
-            parentModule="assessments",
-        )
-        # create assessment child issue via POST
-        iss = api.post(
-            f'{app.config["domain"]}/api/issues',
-            json=issue_data.dict(),
-        )
-        # log issue creation result
-        if iss.ok:
-            logger.info("Issue created successfully.")
-        else:
-            logger.info("Issue was not created.")
+    with create_progress_object() as progress:
+        task = progress.add_task("Creating/updating issue(s) in RegScale...", total=len(sonar_data))
+        for vulnerability in sonar_data:
+            # create issue model
+            issue_data = Issue(
+                title="Sonarcloud Code Scan",  # Required
+                dateCreated=get_current_datetime("%Y-%m-%dT%H:%M:%S"),
+                description=vulnerability["message"],
+                severityLevel=Issue.assign_severity(vulnerability["severity"]),  # Required
+                dueDate=Issue.get_due_date(
+                    severity=vulnerability["severity"].lower(), config=app.config, key="sonarcloud"
+                ),
+                identification="Code scan assessment",
+                status="Open",
+                assessmentId=assessment_id,
+                parentId=parent_id or assessment_id,
+                parentModule=parent_module or "assessments",
+                sourceReport="SonarCloud/Qube",
+                otherIdentifier=vulnerability["key"],
+            )
+            # log issue creation result
+            if issue_data.create_or_update(bulk_create=True, bulk_update=True):
+                logger.debug("Issue was created/updated successfully")
+            else:
+                logger.debug("Issue was not created.")
+            progress.advance(task)
+        Issue.bulk_save(progress)
 
 
 @click.group()
@@ -249,11 +270,29 @@ def sonarcloud() -> None:
 @sonarcloud.command(name="sync_alerts")
 @regscale_id(required=False, default=None)
 @regscale_module(required=False, default=None)
-@click.option("--branch", help="Branch name to filter results, defaults to None")
+@click.option(
+    "--organization",
+    "-o",
+    type=click.STRING,
+    help="Organization name to filter results, defaults to None",
+    default=None,
+)
+@click.option("--branch", "-b", type=click.STRING, help="Branch name to filter results, defaults to None", default=None)
+@click.option("--project_key", "-p", type=click.STRING, help="SonarCloud Project Key, defaults to None", default=None)
 def create_alerts(
-    regscale_id: Optional[int] = None, regscale_module: Optional[str] = None, branch: Optional[str] = None
+    regscale_id: Optional[int] = None,
+    regscale_module: Optional[str] = None,
+    organization: Optional[str] = None,
+    branch: Optional[str] = None,
+    project_key: Optional[str] = None,
 ) -> None:
     """
     Create a child assessment and child issues in RegScale from SonarCloud alerts.
     """
-    create_alert_issues(regscale_id, regscale_module, branch)
+    create_alert_issues(
+        parent_id=regscale_id,
+        parent_module=regscale_module,
+        organization=organization,
+        branch=branch,
+        project_key=project_key,
+    )

regscale/integrations/commercial/tenablev2/jsonl_scanner.py

@@ -446,7 +446,7 @@ class TenableSCJsonlScanner(JSONLScannerIntegration):
             # If no findings were created, return a basic finding
             # Get the IP from the vulnerability directly rather than using passed asset_identifier
             finding_asset_id = vuln.ip or asset_identifier
-
+            logger.debug(item)
             return IntegrationFinding(
                 title=item.get("pluginName", "Unknown Finding"),
                 description=item.get("description", "No description available"),
@@ -456,6 +456,7 @@ class TenableSCJsonlScanner(JSONLScannerIntegration):
                 category="Vulnerability",
                 scan_date=self.scan_date,
                 plugin_name=item.get("pluginName", UNKNOWN_PLUGIN),
+                control_labels=item.get("controlLabels", []),
             )
 
         except Exception as e:

regscale/integrations/commercial/wizv2/async_client.py

@@ -10,6 +10,7 @@ import anyio
 import httpx
 
 from regscale.core.app.utils.app_utils import error_and_exit
+from regscale.integrations.variables import ScannerVariables
 
 logger = logging.getLogger("regscale")
 
@@ -72,7 +73,9 @@ class AsyncWizGraphQLClient:
             logger.debug(f"Variables: {variables}")
 
             try:
-                async with httpx.AsyncClient(timeout=self.timeout) as client:
+                # Get SSL verify setting from scanner variables config
+                ssl_verify = getattr(ScannerVariables, "sslVerify", True)
+                async with httpx.AsyncClient(timeout=self.timeout, verify=ssl_verify) as client:
                     if progress_callback:
                         progress_callback(task_name, "requesting")
 
@@ -81,7 +84,7 @@ class AsyncWizGraphQLClient:
                     if progress_callback:
                         progress_callback(task_name, "processing")
 
-                    if response.raise_for_status():
+                    if not response.is_success:
                         error_and_exit(
                             f"Received non-200 response from GraphQL API: {response.status_code}: {response.text}"
                         )
@@ -304,6 +307,7 @@ def run_async_queries(
     query_configs: List[Dict[str, Any]],
     progress_tracker: Optional[Any] = None,
     max_concurrent: int = 5,
+    timeout: int = 60,
 ) -> List[Tuple[str, List[Dict[str, Any]], Optional[Exception]]]:
     """
     Convenience function to run async queries from synchronous code.
@@ -313,12 +317,15 @@ def run_async_queries(
     :param List[Dict[str, Any]] query_configs: Query configurations
     :param Optional[Any] progress_tracker: Progress tracker
     :param int max_concurrent: Maximum concurrent requests
+    :param int timeout: Request timeout in seconds
     :return: Query results
     :rtype: List[Tuple[str, List[Dict[str, Any]], Optional[Exception]]]
     """
 
     async def _run():
-        client = AsyncWizGraphQLClient(endpoint=endpoint, headers=headers, max_concurrent=max_concurrent)
+        client = AsyncWizGraphQLClient(
+            endpoint=endpoint, headers=headers, max_concurrent=max_concurrent, timeout=timeout
+        )
         return await client.execute_concurrent_queries(query_configs, progress_tracker)
 
     # Use anyio.run for better compatibility