diff --git a/misc/examples/api_examples.py b/misc/examples/api_examples.py index 14ee48e..64edfb8 100644 --- a/misc/examples/api_examples.py +++ b/misc/examples/api_examples.py @@ -1,10 +1,11 @@ -from typing import Optional - +import json import requests +from typing import Optional +from os import getenv -# Set API key and base URL -api_key = "" -url_base = "http://0.0.0.0:8080/api/strelka" +# Environment variables for configuration +api_base = getenv("STRELKA_API_BASE", "http://localhost:8081/api/strelka") +api_key = getenv("STRELKA_API_KEY", "") # Define headers headers = { @@ -13,83 +14,53 @@ "Accept": "application/json", } - -def get_scans_list(page: int = 1, per_page: int = 10) -> Optional[str]: - """ - Get list of scans from the Strelka UI API. - - Args: - page (int): The page number to retrieve. Defaults to 1. - per_page (int): The number of items per page. Defaults to 10. - - Returns: - Optional[str]: The list of scans as a JSON string if the request is successful, None otherwise. - """ - url_route = f"/scans?page={page}&per_page={per_page}" +# Helper function to handle requests +def make_request(method: str, url: str, **kwargs) -> Optional[str]: try: - response = requests.get(url_base + url_route, headers=headers) + response = requests.request(method, url, headers=headers, **kwargs) response.raise_for_status() return response.text + except requests.exceptions.HTTPError as e: + print(f"HTTP Error: {e.response.status_code} for {url}: {e.response.text}") except requests.exceptions.RequestException as e: - print(f"Error: {e}") - return None + print(f"Request Error: {e}") + return None -def get_scan_by_id(scan_id: str) -> Optional[str]: - """ - Get scan details by ID from the Strelka UI API. +# Endpoints +def get_scans_list(page: int = 1, per_page: int = 10) -> Optional[str]: + return make_request("GET", f"{api_base}/scans?page={page}&per_page={per_page}") - Args: - scan_id (str): The ID of the scan to retrieve. - Returns: - Optional[str]: The scan details as a JSON string if the request is successful, None otherwise. - """ - url_route = f"/scans/{scan_id}" - try: - response = requests.get(url_base + url_route, headers=headers) - response.raise_for_status() - return response.text - except requests.exceptions.RequestException as e: - print(f"Error: {e}") - return None +def get_scan_by_id(scan_id: str) -> Optional[str]: + return make_request("GET", f"{api_base}/scans/{scan_id}") def upload_file(filename: str, description: str) -> Optional[str]: - """ - Upload a file to the Strelka UI API. + with open(filename, "rb") as f: + files = {"file": (filename, f)} + data = {"description": description} + return make_request("POST", f"{api_base}/upload", files=files, data=data) - Args: - filename (str): The file path of the file to upload. - description (str): The description of the file. - Returns: - Optional[str]: The response as a JSON string if the upload is successful, None otherwise. +def upload_file_to_vt(file_hash: str, description: str) -> Optional[str]: """ - url_route = "/upload" - upload_headers = { - "X-API-KEY": api_key, + Sample Response: + {"file_id":"07538964-b156-4ed3-93b4-85bcd07b5fbc", + "results":[{"enrichment":{"virustotal":-1},"file":{"depth":0,"flavors":{"mime":["application/zip"],"yara":["encrypted_zip","zip_file"]}... } - - try: - with open(filename, "rb") as f: - file_data = f.read() - - files = [("file", (filename, file_data))] - data = {"description": description} - - response = requests.post( - url_base + url_route, files=files, data=data, headers=upload_headers - ) - response.raise_for_status() - return response.text - except requests.exceptions.RequestException as e: - print(f"Error: {e}") - return None + """ + payload = json.dumps({"description": description, "hash": file_hash}) + return make_request("POST", f"{api_base}/upload", data=payload) if __name__ == "__main__": # Example usage - print(get_scans_list()) - print(get_scan_by_id("")) - print(upload_file("", "This is a test file")) + # print(get_scans_list()) + # print(get_scan_by_id("")) + # print(upload_file("", "This is a test file")) + print( + upload_file_to_vt( + "5da8c98136d98dfec4716edd79c7145f", "VirusTotal Upload - Calc.exe" + ) + ) \ No newline at end of file