Merge pull request #17 from Contrast-Security-OSS/TS-39583

TS-39583: Add API call to mark remediation as failed

Author: JacobMagesHaskinsContrast

src/contrast_api.py

@@ -21,9 +21,22 @@
 import json
 import sys
 from typing import Optional
+from enum import Enum
 import config
 from utils import debug_print
 
+# Define failure categories as an enum to ensure consistency
+class FailureCategory(Enum):
+    INITIAL_BUILD_FAILURE = "INITIAL_BUILD_FAILURE"
+    EXCEEDED_QA_ATTEMPTS = "EXCEEDED_QA_ATTEMPTS"
+    QA_AGENT_FAILURE = "QA_AGENT_FAILURE"
+    GIT_COMMAND_FAILURE = "GIT_COMMAND_FAILURE"
+    AGENT_FAILURE = "AGENT_FAILURE"
+    GENERATE_PR_FAILURE = "GENERATE_PR_FAILURE"
+    HANDLE_PR_MERGE_FAILURE = "HANDLE_PR_MERGE_FAILURE"
+    HANDLE_PR_CLOSE_FAILURE = "HANDLE_PR_CLOSE_FAILURE"
+    GENERAL_FAILURE = "GENERAL_FAILURE"
+
 def normalize_host(host: str) -> str:
     """Remove any protocol prefix from host to prevent double prefixing when constructing URLs."""
     return host.replace('https://', '').replace('http://', '')
@@ -303,3 +316,66 @@ def notify_remediation_pr_closed(remediation_id: str, contrast_host: str, contra
     except json.JSONDecodeError:
         print(f"Error decoding JSON response when notifying Remediation service about closed PR for remediation {remediation_id}.", file=sys.stderr)
         return False
+        
+def notify_remediation_failed(remediation_id: str, failure_category: str, contrast_host: str, contrast_org_id: str, contrast_app_id: str, contrast_auth_key: str, contrast_api_key: str) -> bool:
+    """Notifies the Remediation backend service that a remediation has failed.
+
+    Args:
+        remediation_id: The ID of the remediation.
+        failure_category: The category of failure (e.g., "INITIAL_BUILD_FAILURE").
+        contrast_host: The Contrast Security host URL.
+        contrast_org_id: The organization ID.
+        contrast_app_id: The application ID.
+        contrast_auth_key: The Contrast authorization key.
+        contrast_api_key: The Contrast API key.
+
+    Returns:
+        bool: True if the notification was successful, False otherwise.
+    """
+    debug_print(f"--- Notifying Remediation service about failed remediation {remediation_id} with category {failure_category} ---")
+    api_url = f"https://{normalize_host(contrast_host)}/api/v4/aiml-remediation/organizations/{contrast_org_id}/applications/{contrast_app_id}/remediations/{remediation_id}/failed"
+
+    headers = {
+        "Authorization": contrast_auth_key,
+        "API-Key": contrast_api_key,
+        "Content-Type": "application/json",
+        "Accept": "application/json",
+        "User-Agent": config.USER_AGENT
+    }
+    
+    payload = {
+        "failureCategory": failure_category
+    }
+
+    try:
+        debug_print(f"Making PUT request to: {api_url}")
+        debug_print(f"Payload: {json.dumps(payload)}") 
+        response = requests.put(api_url, headers=headers, json=payload)
+        response.raise_for_status()  # Raises an HTTPError for bad responses (4xx or 5xx)
+
+        debug_print(f"Remediation failed notification API response status code: {response.status_code}")
+        
+        if response.status_code == 204:
+            debug_print(f"Successfully notified Remediation service API about failed remediation {remediation_id}")
+            return True
+        else:
+            error_message = "Unknown error"
+            try:
+                response_json = response.json()
+                if "messages" in response_json and response_json["messages"]:
+                    error_message = response_json["messages"][0]
+            except:
+                error_message = response.text
+                
+            print(f"Failed to notify Remediation service about failed remediation {remediation_id}. Error: {error_message}", file=sys.stderr)
+            return False
+
+    except requests.exceptions.HTTPError as e:
+        print(f"HTTP error notifying Remediation service about failed remediation {remediation_id}: {e.response.status_code} - {e.response.text}", file=sys.stderr)
+        return False
+    except requests.exceptions.RequestException as e:
+        print(f"Request error notifying Remediation service about failed remediation {remediation_id}: {e}", file=sys.stderr)
+        return False
+    except json.JSONDecodeError:
+        print(f"Error decoding JSON response when notifying Remediation service about failed remediation {remediation_id}.", file=sys.stderr)
+        return False

src/main.py

@@ -160,6 +160,23 @@ def main():
             error_analysis = extract_build_errors(prefix_build_output)
             print("\n❌ Build is broken ❌ -- No fix attempted.")
             print(f"Build output:\n{error_analysis}")
+            
+            # Notify the Remediation service about the failed build
+            remediation_notified = contrast_api.notify_remediation_failed(
+                remediation_id=remediation_id,
+                failure_category=contrast_api.FailureCategory.INITIAL_BUILD_FAILURE.value,
+                contrast_host=config.CONTRAST_HOST,
+                contrast_org_id=config.CONTRAST_ORG_ID,
+                contrast_app_id=config.CONTRAST_APP_ID,
+                contrast_auth_key=config.CONTRAST_AUTHORIZATION_KEY,
+                contrast_api_key=config.CONTRAST_API_KEY
+            )
+            
+            if remediation_notified:
+                print(f"Successfully notified Remediation service about failed build for remediation {remediation_id}.", flush=True)
+            else:
+                print(f"Warning: Failed to notify Remediation service about failed build for remediation {remediation_id}.", flush=True)
+                
             git_handler.cleanup_branch(new_branch_name)
             sys.exit(1) # Exit if the build is broken, no point in proceeding
 
@@ -221,10 +238,34 @@ def main():
                 # Skip PR creation if QA was run and the build is failing
                 # or if the QA agent encountered an error (detected by checking qa_summary_log entries)
                 if (used_build_command and not build_success) or any(s.startswith("Error during QA agent execution:") for s in qa_summary_log):
+                    failure_category = ""
+                    
                     if any(s.startswith("Error during QA agent execution:") for s in qa_summary_log):
                         print("\n--- Skipping PR creation as QA Agent encountered an error ---")
+                        failure_category = contrast_api.FailureCategory.QA_AGENT_FAILURE.value
                     else:
                         print("\n--- Skipping PR creation as QA Agent failed to fix build issues ---")
+                        # Check if we've exhausted all retry attempts
+                        if len(qa_summary_log) >= max_qa_attempts_setting:
+                            failure_category = contrast_api.FailureCategory.EXCEEDED_QA_ATTEMPTS.value
+                    
+                    # Notify the Remediation service about the failed remediation if we have a failure category
+                    if failure_category:
+                        remediation_notified = contrast_api.notify_remediation_failed(
+                            remediation_id=remediation_id,
+                            failure_category=failure_category,
+                            contrast_host=config.CONTRAST_HOST,
+                            contrast_org_id=config.CONTRAST_ORG_ID,
+                            contrast_app_id=config.CONTRAST_APP_ID,
+                            contrast_auth_key=config.CONTRAST_AUTHORIZATION_KEY,
+                            contrast_api_key=config.CONTRAST_API_KEY
+                        )
+                        
+                        if remediation_notified:
+                            print(f"Successfully notified Remediation service about {failure_category} for remediation {remediation_id}.", flush=True)
+                        else:
+                            print(f"Warning: Failed to notify Remediation service about {failure_category} for remediation {remediation_id}.", flush=True)
+                    
                     git_handler.cleanup_branch(new_branch_name)
                     continue # Move to the next vulnerability