#!/usr/bin/env python3
"""
Test script for the new policy results CSV export API.
This script demonstrates how to authenticate and call the endpoint with different filters.

How to use with an API key:
1. Get your API key from the web UI or administrator
2. Run this script with: python3 test_policy_csv.py --url http://127.0.0.1:5003 --token "your_api_key_here" --output tmp/test.csv
"""

import requests
import json
import argparse
import os
from urllib.parse import urljoin

def get_token_from_api_key(base_url, api_key):
    """Get a JWT token by authenticating with an API key"""
    auth_url = urljoin(base_url, "/v1/login")
    
    try:
        # The login endpoint expects the API key in the JSON payload as 'api_token'
        payload = {"api_token": api_key}
        response = requests.post(auth_url, json=payload)
        
        if response.status_code != 200:
            print(f"Authentication failed with API key: {response.text}")
            return None
        
        data = response.json()
        jwt_token = data.get("access_token")
        if jwt_token:
            print("Successfully authenticated with API key")
            return jwt_token
        else:
            print(f"API key authentication succeeded but no token returned: {data}")
            return None
    except Exception as e:
        print(f"Error authenticating with API key: {str(e)}")
        return None

def get_token_from_credentials(base_url, username, password):
    """Get a JWT token by authenticating with username/password"""
    auth_url = urljoin(base_url, "/v1/login")
    
    try:
        response = requests.post(
            auth_url, 
            json={"username": username, "password": password}
        )
        
        if response.status_code != 200:
            print(f"Authentication failed with credentials: {response.text}")
            return None
        
        data = response.json()
        jwt_token = data.get("access_token")
        if jwt_token:
            print("Successfully authenticated with username/password")
            return jwt_token
        else:
            print(f"Credential authentication succeeded but no token returned: {data}")
            return None
    except Exception as e:
        print(f"Error authenticating with credentials: {str(e)}")
        return None

def download_policy_csv(base_url, jwt_token, filters=None, output_file=None):
    """
    Download policy results as CSV with optional filters
    
    Args:
        base_url: API base URL
        jwt_token: JWT token for authentication
        filters: Dict with optional filters (rule_id, product, result_kind, result_level, filename)
        output_file: Path to save the CSV file (if None, prints to console)
    """
    # Build URL with filters
    csv_url = urljoin(base_url, "/evidence/policy-results/csv")
    params = filters or {}
    
    # Make the request
    # Remove 'Bearer ' prefix if it's already included in the token
    if jwt_token.startswith('Bearer '):
        auth_header = jwt_token
    else:
        auth_header = f"Bearer {jwt_token}"
    
    headers = {"Authorization": auth_header}
    print(f"Requesting: {csv_url}")
    
    try:
        response = requests.get(csv_url, headers=headers, params=params)
        
        if response.status_code != 200:
            print(f"Download failed (Status code: {response.status_code}): {response.text}")
            return False
        
        # Get the filename from Content-Disposition header if available
        content_disposition = response.headers.get('Content-Disposition', '')
        filename = "policy_results.csv"  # Default name
        
        if 'filename=' in content_disposition:
            filename = content_disposition.split('filename=')[1].strip('"')
        
        # Write to file or print to console
        if output_file:
            # Create directory if it doesn't exist
            output_dir = os.path.dirname(output_file)
            if output_dir and not os.path.exists(output_dir):
                os.makedirs(output_dir)
                
            with open(output_file, 'wb') as f:
                f.write(response.content)
            print(f"CSV saved to {output_file}")
        else:
            print(response.text)
        
        return True
    except Exception as e:
        print(f"Error downloading CSV: {str(e)}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Download policy results as CSV')
    parser.add_argument('--url', help='Base URL of the API server', default='http://localhost:5000')
    parser.add_argument('--username', help='Username for authentication')
    parser.add_argument('--password', help='Password for authentication')
    parser.add_argument('--token', help='API key for authentication')
    parser.add_argument('--jwt', help='Pre-existing JWT token (if provided, no authentication needed)')
    parser.add_argument('--rule-id', help='Filter by rule ID')
    parser.add_argument('--product', help='Filter by product key')
    parser.add_argument('--result-kind', help='Filter by result kind')
    parser.add_argument('--result-level', help='Filter by result level')
    parser.add_argument('--filename', help='Custom filename for the CSV')
    parser.add_argument('--output', help='Path to save the CSV file (if not specified, print to console)')
    
    args = parser.parse_args()
    
    # First, check if we already have a JWT token
    jwt_token = args.jwt
    
    # If no JWT, try to authenticate with API key
    if not jwt_token and args.token:
        jwt_token = get_token_from_api_key(args.url, args.token)
    
    # If still no JWT, try username/password
    if not jwt_token and (args.username and args.password):
        jwt_token = get_token_from_credentials(args.url, args.username, args.password)
    
    if not jwt_token:
        print("Error: No valid authentication method provided. Please use --token, --jwt, or --username/--password.")
        return 1
    
    # Build filters from arguments
    filters = {}
    if args.rule_id:
        filters['rule_id'] = args.rule_id
    if args.product:
        filters['product'] = args.product
    if args.result_kind:
        filters['result_kind'] = args.result_kind
    if args.result_level:
        filters['result_level'] = args.result_level
    if args.filename:
        filters['filename'] = args.filename
    
    # Download CSV using the JWT token
    print(f"Downloading policy results CSV to {args.output or 'console'}")
    success = download_policy_csv(args.url, jwt_token, filters, args.output)
    return 0 if success else 1

if __name__ == "__main__":
    exit(main())