Spaces:
Sleeping
Sleeping
| import argparse | |
| import os | |
| import logging | |
| from .client import ScribeClient | |
| import json | |
| from dotenv import load_dotenv | |
| dotenv_path = os.path.join(os.path.dirname(__file__), "mcp_app", "env") # relative to cli.py location | |
| load_dotenv(dotenv_path) | |
| def main(): | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") | |
| parser = argparse.ArgumentParser(description="Scribe Client CLI") | |
| parser.add_argument("--api-call", choices=[ | |
| "get-products", | |
| "get-datasets", | |
| "get-policy-results", | |
| "list-attestations", | |
| "get-attestation", | |
| "delete-attestation", | |
| "get-product-vulnerabilities", | |
| "get-latest-attestation", | |
| "query-vulnerabilities", | |
| "query-products", | |
| "query-policy-results", | |
| "query-lineage", | |
| "query-risk", | |
| "query-findings", | |
| "get-dataset-links", | |
| "get-dataset-data", | |
| "get-dataset-tables-dict", | |
| "get-dataset-tables-md", | |
| ], help="Which API call to execute") | |
| parser.add_argument("--api-token", required=False, default=os.getenv('SCRIBE_TOKEN'), help="Your API token from ScribeHub integrations page (defaults to SCRIBE_TOKEN environment variable)") | |
| parser.add_argument("--env", default="prod", choices=["prod", "dev", "test", "ci"], help="Which environment to use") | |
| parser.add_argument("--product-name", required=False, help="The name of the product (required for specific API calls)") | |
| parser.add_argument("--attestation-id", required=False, help="The ID of the attestation (required for get-attestation API call)") | |
| parser.add_argument("--query", required=False, help="The query to run on the dataset (required for query commands)") | |
| parser.add_argument("--lineage-graph-file", required=False, help="The file to save the lineage graph to (optional)") | |
| parser.add_argument("--username", required=False, default = os.getenv('SUPERSET_DEV_USERNAME'), help="Superset username (required for Superset API calls)") | |
| parser.add_argument("--password", required=False, default = os.getenv('SUPERSET_DEV_PASSWORD'), help="Superset password (required for Superset API calls)") | |
| parser.add_argument("--mcp", action="store_true", help="Run the MCP FastAPI server (requires [mcp] extras)") | |
| args = parser.parse_args() | |
| if args.mcp: | |
| import traceback | |
| try: | |
| from scribe_python_client.mcp_app.main import run_mcp_server | |
| except Exception as e: | |
| logging.error("FastAPI server not available. Install with: pip install scribe-python-client[mcp]") | |
| traceback.print_exc() | |
| exit(1) | |
| run_mcp_server() | |
| return | |
| if not args.api_token: | |
| logging.error("No API token provided. Please set the SCRIBE_TOKEN environment variable or pass --api-token.") | |
| exit(1) | |
| if args.api_call in ["get-policy-results", "get-product-vulnerabilities", "get-latest-attestation"] and not args.product_name: | |
| logging.error("The --product-name argument is required for the selected API call.") | |
| exit(1) | |
| if args.api_call in ["get-attestation", "delete-attestation"] and not args.attestation_id: | |
| logging.error("The --attestation-id argument is required for the selected API call.") | |
| exit(1) | |
| if args.api_call in ["query-vulnerabilities", "query-products", "query-policy-results", "query-lineage"] and not args.query: | |
| logging.error("The --query argument is required for the selected API call.") | |
| exit(1) | |
| base_url = "https://api.scribesecurity.com" | |
| if args.env != "prod": | |
| base_url = f"https://api.{args.env}.scribesecurity.com" | |
| if args.api_call in ["get-dataset-links", "get-dataset-data", "get-dataset-tables-dict", "get-dataset-tables-md"]: | |
| if not args.username or not args.password: | |
| logging.error("--username and --password are required for Superset table export commands.") | |
| exit(1) | |
| client = ScribeClient(api_token=args.api_token, base_url=base_url, env = args.env, username= args.username, password=args.password) | |
| # lineage_prompt = client.get_dataset_prompt(client.scribe_lineage_dataset) | |
| # print(f"Lineage prompt:\n {lineage_prompt}") | |
| logging.info(f"Executing API call: {args.api_call}") | |
| if args.api_call == "get-products": | |
| print(client.get_products_str()) | |
| elif args.api_call == "get-datasets": | |
| print(client.get_datasets()) | |
| elif args.api_call == "get-policy-results": | |
| print(client.get_policy_results_str(logical_app=args.product_name)) | |
| elif args.api_call == "list-attestations": | |
| print(client.list_attestations()) | |
| elif args.api_call == "get-attestation": | |
| print(client.get_attestation(args.attestation_id)) | |
| elif args.api_call == "delete-attestation": | |
| success = client.delete_attestation(args.attestation_id) | |
| if success: | |
| print(f"Attestation {args.attestation_id} deleted successfully.") | |
| else: | |
| print(f"Failed to delete attestation {args.attestation_id}.") | |
| elif args.api_call == "get-product-vulnerabilities": | |
| print(client.get_product_vulnerabilities_str(logical_app=args.product_name)) | |
| elif args.api_call == "get-latest-attestation": | |
| criteria = {"name": args.product_name} | |
| print(client.get_latest_attestation(criteria=criteria)) | |
| elif args.api_call == "query-vulnerabilities": | |
| print(client.query_vulnerabilities(querystr=args.query)) | |
| elif args.api_call == "query-products": | |
| print(client.query_products(querystr=args.query)) | |
| elif args.api_call == "query-policy-results": | |
| print(client.query_policy_results(querystr=args.query)) | |
| elif args.api_call == "query-lineage": | |
| print(client.query_lineage(querystr=args.query)) | |
| if args.lineage_graph_file: | |
| client.query_lineage(args.query, graph_filename=args.lineage_graph_file) | |
| # logging.info(f"Lineage graph saved to {args.lineage_graph_file}") | |
| elif args.api_call == "query-risk": | |
| print(client.query_risk(querystr=args.query)) | |
| elif args.api_call == "query-findings": | |
| print(client.query_findings(querystr=args.query)) | |
| # Get dataset metadata section: | |
| elif args.api_call == "get-dataset-links": | |
| links = client.get_dataset_links() | |
| if links: | |
| print("Dataset Links:") | |
| for link in links: | |
| print(link) | |
| else: | |
| logging.error("No dataset links found.") | |
| result = None | |
| elif args.api_call == "get-dataset-data": | |
| dataset_links = client.get_dataset_links() | |
| if dataset_links: | |
| client.get_dataset_data(dataset_links, output_file="dataset_metadata.csv") | |
| logging.info("Metadata export completed successfully.") | |
| else: | |
| logging.error("No dataset links found.") | |
| result = None | |
| elif args.api_call == "get-dataset-tables-dict": | |
| dataset_links = client.get_dataset_links() | |
| if dataset_links: | |
| client.get_dataset_data(dataset_links, output_file="dataset_metadata.csv") | |
| logging.info("Metadata export completed successfully.") | |
| else: | |
| logging.error("No dataset links found.") | |
| result = None | |
| tables_dict = client.get_dataset_tables_dict() | |
| #Output path is set to docs/dataset_tables.json | |
| output_path = os.path.join("docs", "dataset_tables.json") | |
| with open(output_path, "w", encoding="utf-8") as f: | |
| json.dump(tables_dict, f, indent=2) | |
| logging.info(f"Saved table metadata dictionary to {output_path}") | |
| if result is not None: | |
| print(result) | |
| elif args.api_call == "get-dataset-tables-md": | |
| if not client.jwt_token: | |
| username = input("Enter your Superset username: ") | |
| password = input("Enter your Superset password: ") | |
| client.username = username | |
| client.password = password | |
| if not client.authenticate(): | |
| logging.error("Superset login failed.") | |
| exit(1) | |
| dataset_links = client.get_dataset_links() | |
| if dataset_links: | |
| client.get_dataset_data(dataset_links, output_file="dataset_metadata.csv") | |
| md_output = client.get_dataset_tables_md() | |
| #Output path is set to docs/dataset_tables.md | |
| output_path = os.path.join("docs", "dataset_tables.md") | |
| with open(output_path, "w", encoding="utf-8") as f: | |
| f.write(md_output) | |
| logging.info(f"Saved table metadata Markdown to {output_path}") | |
| else: | |
| logging.error("No dataset links found.") | |
| if __name__ == "__main__": | |
| main() |