runzero_api_client.py

#!/usr/bin/env python3
"""
runZero API Client
Pulls findings and vulnerabilities data from runZero API and aggregates by severity.

Documentation: https://help.runzero.com/docs/leveraging-the-api/
"""

import os
import sys
import json
import requests
from collections import defaultdict
from typing import Dict, List, Optional
import argparse


class RunZeroClient:
    """Client for interacting with the runZero API"""

    BASE_URL = "https://console.runzero.com/api/v1.0"

    def __init__(self, api_token: Optional[str] = None, client_id: Optional[str] = None,
                 client_secret: Optional[str] = None):
        """
        Initialize the runZero API client.

        Args:
            api_token: Export/Organization/Account API token (Bearer token)
            client_id: OAuth2 client ID (alternative to api_token)
            client_secret: OAuth2 client secret (required with client_id)
        """
        self.session = requests.Session()

        if api_token:
            self.session.headers.update({
                'Authorization': f'Bearer {api_token}',
                'Content-Type': 'application/json'
            })
        elif client_id and client_secret:
            # Get OAuth2 token
            access_token = self._get_oauth_token(client_id, client_secret)
            self.session.headers.update({
                'Authorization': f'Bearer {access_token}',
                'Content-Type': 'application/json'
            })
        else:
            raise ValueError("Either api_token or (client_id and client_secret) must be provided")

    def _get_oauth_token(self, client_id: str, client_secret: str) -> str:
        """
        Get OAuth2 access token using client credentials.

        Args:
            client_id: OAuth2 client ID
            client_secret: OAuth2 client secret

        Returns:
            Access token string
        """
        token_url = f"{self.BASE_URL}/account/api/token"
        data = {
            'grant_type': 'client_credentials',
            'client_id': client_id,
            'client_secret': client_secret
        }

        response = requests.post(
            token_url,
            data=data,
            headers={'Content-Type': 'application/x-www-form-urlencoded'}
        )
        response.raise_for_status()

        return response.json()['access_token']

    def get_findings(self, search: Optional[str] = None) -> List[Dict]:
        """
        Get findings from runZero.

        Args:
            search: Optional search query to filter findings

        Returns:
            List of finding dictionaries
        """
        url = f"{self.BASE_URL}/export/org/findings.jsonl"
        params = {}
        if search:
            params['search'] = search

        response = self.session.get(url, params=params)
        response.raise_for_status()

        # Parse JSONL format (one JSON object per line)
        findings = []
        for line in response.text.strip().split('\n'):
            if line:
                findings.append(json.loads(line))

        return findings

    def get_vulnerabilities(self, search: Optional[str] = None) -> List[Dict]:
        """
        Get vulnerabilities from runZero.

        Args:
            search: Optional search query to filter vulnerabilities

        Returns:
            List of vulnerability dictionaries
        """
        url = f"{self.BASE_URL}/export/org/vulnerabilities.json"
        params = {}
        if search:
            params['search'] = search

        response = self.session.get(url, params=params)
        response.raise_for_status()

        return response.json()

    def aggregate_findings_by_severity(self, findings: List[Dict]) -> Dict[str, int]:
        """
        Aggregate findings by severity level.

        Args:
            findings: List of findings

        Returns:
            Dictionary mapping severity to count
        """
        severity_counts = defaultdict(int)

        for finding in findings:
            severity = finding.get('severity', 'unknown').lower()
            severity_counts[severity] += 1

        return dict(severity_counts)

    def aggregate_vulnerabilities_by_severity(self, vulnerabilities: List[Dict]) -> Dict[str, int]:
        """
        Aggregate vulnerabilities by severity level.

        Args:
            vulnerabilities: List of vulnerabilities

        Returns:
            Dictionary mapping severity to count
        """
        severity_counts = defaultdict(int)

        for vuln in vulnerabilities:
            severity = vuln.get('severity', 'unknown').lower()
            severity_counts[severity] += 1

        return dict(severity_counts)


def print_severity_summary(title: str, severity_counts: Dict[str, int]):
    """Print a formatted summary of severity counts."""
    print(f"\n{title}")
    print("=" * 50)

    # Sort by severity priority
    severity_order = ['critical', 'high', 'medium', 'low', 'info', 'unknown']
    total = 0

    for severity in severity_order:
        count = severity_counts.get(severity, 0)
        if count > 0:
            print(f"{severity.capitalize():12s}: {count:,}")
            total += count

    # Add any other severities not in the standard order
    for severity, count in severity_counts.items():
        if severity not in severity_order:
            print(f"{severity.capitalize():12s}: {count:,}")
            total += count

    print("-" * 50)
    print(f"{'Total':12s}: {total:,}")


def main():
    parser = argparse.ArgumentParser(
        description='Pull findings and vulnerabilities data from runZero API',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Using API token from environment variable
  export RUNZERO_API_TOKEN="your-token-here"
  python runzero_api_client.py

  # Using API token from command line
  python runzero_api_client.py --api-token "your-token-here"

  # Using OAuth2 credentials
  python runzero_api_client.py --client-id "id" --client-secret "secret"

  # Filter findings by severity
  python runzero_api_client.py --search-findings "severity:critical"

  # Export to JSON file
  python runzero_api_client.py --output results.json
        """
    )

    # Authentication options
    auth_group = parser.add_argument_group('Authentication')
    auth_group.add_argument(
        '--api-token',
        help='runZero API token (or set RUNZERO_API_TOKEN env var)',
        default=os.getenv('RUNZERO_API_TOKEN')
    )
    auth_group.add_argument(
        '--client-id',
        help='OAuth2 client ID (or set RUNZERO_CLIENT_ID env var)',
        default=os.getenv('RUNZERO_CLIENT_ID')
    )
    auth_group.add_argument(
        '--client-secret',
        help='OAuth2 client secret (or set RUNZERO_CLIENT_SECRET env var)',
        default=os.getenv('RUNZERO_CLIENT_SECRET')
    )

    # Query options
    query_group = parser.add_argument_group('Query Options')
    query_group.add_argument(
        '--search-findings',
        help='Search query to filter findings (e.g., "severity:critical")',
        default=None
    )
    query_group.add_argument(
        '--search-vulns',
        help='Search query to filter vulnerabilities',
        default=None
    )

    # Output options
    output_group = parser.add_argument_group('Output Options')
    output_group.add_argument(
        '--output', '-o',
        help='Output file path for JSON results (optional)',
        default=None
    )
    output_group.add_argument(
        '--quiet', '-q',
        action='store_true',
        help='Suppress console output (only save to file)'
    )

    args = parser.parse_args()

    # Validate authentication
    if not args.api_token and not (args.client_id and args.client_secret):
        parser.error("Either --api-token or both --client-id and --client-secret must be provided")

    try:
        # Initialize client
        if args.api_token:
            client = RunZeroClient(api_token=args.api_token)
        else:
            client = RunZeroClient(client_id=args.client_id, client_secret=args.client_secret)

        # Fetch findings
        if not args.quiet:
            print("Fetching findings from runZero...")
        findings = client.get_findings(search=args.search_findings)
        findings_by_severity = client.aggregate_findings_by_severity(findings)

        # Fetch vulnerabilities
        if not args.quiet:
            print("Fetching vulnerabilities from runZero...")
        vulnerabilities = client.get_vulnerabilities(search=args.search_vulns)
        vulns_by_severity = client.aggregate_vulnerabilities_by_severity(vulnerabilities)

        # Prepare results
        results = {
            'findings': {
                'total': len(findings),
                'by_severity': findings_by_severity,
                'raw_data': findings if args.output else None
            },
            'vulnerabilities': {
                'total': len(vulnerabilities),
                'by_severity': vulns_by_severity,
                'raw_data': vulnerabilities if args.output else None
            }
        }

        # Print summary to console
        if not args.quiet:
            print_severity_summary("FINDINGS BY SEVERITY", findings_by_severity)
            print_severity_summary("\nVULNERABILITIES BY SEVERITY", vulns_by_severity)

        # Save to file if requested
        if args.output:
            with open(args.output, 'w') as f:
                json.dump(results, f, indent=2)
            if not args.quiet:
                print(f"\n\nResults saved to: {args.output}")

        return 0

    except requests.exceptions.HTTPError as e:
        print(f"HTTP Error: {e}", file=sys.stderr)
        print(f"Response: {e.response.text}", file=sys.stderr)
        return 1
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        return 1


if __name__ == '__main__':
    sys.exit(main())