diff --git a/get_results.py b/get_results.py index 654808c..901d9da 100644 --- a/get_results.py +++ b/get_results.py @@ -4,15 +4,20 @@ import pandas as pd from dotenv import load_dotenv import csv +import argparse def process_response_data(request_id, response_data, results_df): - for item in response_data.get('data', []): - overall_status = item["resultsSummary"].get('status', 'N/A') - score = item["resultsSummary"].get('metadata', {}).get('finalScore', 'N/A') + data = response_data.get("mediaList") or response_data.get("data") or [] + for item in data: + results_summary = item.get("resultsSummary", {}) + + overall_status = results_summary.get('status', 'UNABLE_TO_EVALUATE') if results_summary else 'UNABLE_TO_EVALUATE' + score = results_summary.get('metadata', {}).get('finalScore', '') if results_summary else '' # Append the data to the DataFrame results_df = pd.concat([results_df, pd.DataFrame([{ - 'request_id': request_id, + 'file_name': item.get('originalFileName', ''), + 'request_id': request_id if request_id else item.get('requestId', ''), 'status': overall_status, 'score': score }])], ignore_index=True) @@ -27,30 +32,36 @@ def fetch_data_from_api(url, headers): print(f"An error occurred: {e}") return None -def get_media_detail(request_id, token, results_df): - page_index = 1 +def get_media_detail(request_id, token, results_df, include_all_users=False): + page_index = 0 headers = {"x-api-key": token, "Content-Type": "application/json"} - url = f"https://api.prd.realitydefender.xyz/api/media/users/{request_id}?pageIndex={page_index}" if request_id == "": while True: + url = (f"https://api.dev.realitydefender.xyz/api/v2/media/users/pages/{page_index}?userIds=[]" + if include_all_users + else f"https://api.dev.realitydefender.xyz/api/v2/media/users/pages/{page_index}") print(f"getting page {page_index}") response_data = fetch_data_from_api(url, headers) - if response_data is None or not response_data.get('data'): - print(f"No data found for page {page_index}. Skipping.") - page_index += 1 - continue + if response_data is None: + print(f"Error fetching page {page_index}. Stopping.") + break + + if not response_data.get('mediaList'): + print(f"No data found for page {page_index}. Stopping.") + break results_df = process_response_data(request_id, response_data, results_df) - total_pages = response_data.get('totalPages', 1) - + total_pages = response_data.get('totalPages', 0) + if page_index >= total_pages: break page_index += 1 else: + url = f"https://api.dev.realitydefender.xyz/api/media/users/{request_id}" response_data = fetch_data_from_api(url, headers) print(f"fetching {request_id}") if response_data: @@ -61,21 +72,19 @@ def get_media_detail(request_id, token, results_df): if __name__ == "__main__": load_dotenv() token = os.getenv("RD_API") - results_df = pd.DataFrame(columns=['request_id', 'status', 'score']) + results_df = pd.DataFrame(columns=['file_name', 'request_id', 'status', 'score']) - if len(sys.argv) > 2: - print("Usage: python sample_get_script.py ") - sys.exit(1) - elif len(sys.argv) == 1: - request_id = "" - results_df = get_media_detail(request_id, token, results_df) - results_df.to_csv('results.csv', index=False) - print("Results saved to results.csv") - else: - csv_file_path = sys.argv[1] + # Set up argument parser + parser = argparse.ArgumentParser(description='Fetch media details from Reality Defender API') + parser.add_argument('--get-all', action='store_true', help='Get media details for all users') + parser.add_argument('csv_file', nargs='?', default=None, help='Path to CSV file containing request IDs') + + args = parser.parse_args() + + if args.csv_file: try: # Load the CSV into a DataFrame - df = pd.read_csv(csv_file_path) + df = pd.read_csv(args.csv_file) # Validate that the required columns are present required_columns = {'file_name', 'request_id'} @@ -92,8 +101,13 @@ def get_media_detail(request_id, token, results_df): merged_df.to_csv('results.csv', index=False) print("Results saved to results.csv") except FileNotFoundError: - print(f"Error: File {csv_file_path} not found.") + print(f"Error: File {args.csv_file} not found.") sys.exit(1) except Exception as e: print(f"Error: {e}") sys.exit(1) + else: + request_id = "" + results_df = get_media_detail(request_id, token, results_df, include_all_users=args.get_all) + results_df.to_csv('results.csv', index=False) + print("Results saved to results.csv")