File size: 6,663 Bytes
e17cbcc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | #!/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())
|