vectr.py

import requests
import json
from typing import Dict, List, Optional

class VectrAPIClient:
    """A client to interact with the VECTR GraphQL API for pulling test cases."""
    
    def __init__(self, api_url: str, api_key: str):
        self.api_url = api_url
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

    def execute_query(self, query: str, variables: Optional[Dict] = None) -> Dict:
        """Execute a GraphQL query and return the response."""
        payload = {"query": query}
        if variables:
            payload["variables"] = variables
        
        try:
            response = requests.post(
                self.api_url,
                headers=self.headers,
                json=payload,
                timeout=10
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            raise Exception(f"HTTP error occurred: {http_err}")
        except requests.exceptions.RequestException as req_err:
            raise Exception(f"Request error occurred: {req_err}")
        except ValueError as json_err:
            raise Exception(f"JSON decode error: {json_err}")

    def get_test_cases(self, database: str, campaign_id: str) -> List[Dict]:
        """Retrieve test cases for a given database and campaign."""
        query = """
        query GetTestCases($database: String!, $campaignId: String!) {
            testCases(database: $database, campaignId: $campaignId) {
                id
                name
                description
                phase
                technique
                defenses
                redTools {
                    name
                }
                tags
                organization
            }
        }
        """
        variables = {
            "database": database,
            "campaignId": campaign_id
        }
        
        response = self.execute_query(query, variables)
        
        if "errors" in response:
            raise Exception(f"GraphQL errors: {response['errors']}")
        
        test_cases = response.get("data", {}).get("testCases", [])
        return test_cases

def main():
    # Configuration - Replace with your VECTR instance details
    API_URL = "https://vectr.internal/sra-purpletools-rest/graphql"  # Update with your VECTR GraphQL endpoint
    API_KEY = "PUTYOURAPIKEYHERE"  # Replace with your actual API key
    DATABASE = "MY_USER_DB"  # Replace with your VECTR database name
    CAMPAIGN_ID = "YOUR_CAMPAIGN_ID"  # Replace with your campaign ID

    # Initialize the VECTR API client
    client = VectrAPIClient(api_url=API_URL, api_key=API_KEY)
    
    try:
        # Fetch test cases
        test_cases = client.get_test_cases(database=DATABASE, campaign_id=CAMPAIGN_ID)
        
        # Print test cases
        if not test_cases:
            print("No test cases found.")
            return
        
        print(f"Found {len(test_cases)} test cases:")
        for test_case in test_cases:
            print("\nTest Case:")
            print(f"ID: {test_case['id']}")
            print(f"Name: {test_case['name']}")
            print(f"Description: {test_case['description'] or 'N/A'}")
            print(f"Phase: {test_case['phase'] or 'N/A'}")
            print(f"Technique: {test_case['technique'] or 'N/A'}")
            print(f"Defenses: {', '.join(test_case['defenses']) if test_case['defenses'] else 'None'}")
            print(f"Red Tools: {', '.join(tool['name'] for tool in test_case['redTools']) if test_case['redTools'] else 'None'}")
            print(f"Tags: {', '.join(test_case['tags']) if test_case['tags'] else 'None'}")
            print(f"Organization: {test_case['organization'] or 'N/A'}")
            
    except Exception as e:
        print(f"Error: {str(e)}")

if __name__ == "__main__":
    main()