| |
| """ |
| 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: |
| |
| 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) |
| """ |
| |
| csv_url = urljoin(base_url, "/evidence/policy-results/csv") |
| params = filters or {} |
| |
| |
| |
| 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 |
| |
| |
| content_disposition = response.headers.get('Content-Disposition', '') |
| filename = "policy_results.csv" |
| |
| if 'filename=' in content_disposition: |
| filename = content_disposition.split('filename=')[1].strip('"') |
| |
| |
| if output_file: |
| |
| 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() |
| |
| |
| jwt_token = args.jwt |
| |
| |
| if not jwt_token and args.token: |
| jwt_token = get_token_from_api_key(args.url, args.token) |
| |
| |
| 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 |
| |
| |
| 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 |
| |
| |
| 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()) |
|
|
