Blame for scribe_python_client/cli.py

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",
        "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 == "get-attestation" and not args.attestation_id:
        logging.error("The --attestation-id argument is required for the get-attestation 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 == "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()