| import textwrap |
| from fastapi import FastAPI, HTTPException, Query |
| from fastapi.responses import JSONResponse |
| from scribe_python_client.client import ScribeClient |
| from dotenv import load_dotenv |
| from pydantic import BaseModel |
| from typing import * |
| from pathlib import Path |
| from contextlib import asynccontextmanager |
| import logging |
| import json |
| import os |
|
|
|
|
| load_dotenv() |
| API_TOKEN = os.getenv("SCRIBE_TOKEN") |
| client = ScribeClient(api_token=API_TOKEN) |
|
|
| |
| toolDocsPath = Path(__file__).parent / "toolsDescriptions/tool_docs.json" |
|
|
| class Filter(BaseModel): |
| col: str |
| op: str |
| val: Any |
|
|
| class Metric(BaseModel): |
| label: str |
| expressionType: str |
| sqlExpression: Optional[str] = None |
|
|
| class VulnerabilityQuery(BaseModel): |
| column: Optional[List[str]] = [] |
| filters: Optional[List[Filter]] = [] |
| metrics: Optional[List[Metric]] = [] |
| orderby: Optional[List[List[Any]]] = [] |
| row_limit: Optional[int] = 20 |
|
|
| class QueryResponse(BaseModel): |
| format: str |
| status: bool |
| status_info: str |
| json_data: object |
| str: str |
|
|
|
|
| def get_products_list() -> list: |
| result = client.get_products() |
| listOfProducts = [] |
| for i in range(len(result)): |
| listOfProducts.append(result[i]["name"]) |
| return listOfProducts |
|
|
| app = FastAPI() |
|
|
| @app.post("/generate_query", operation_id="generateQuery") |
| def build_llm_prompt_from_docs( input: str = Query(...), function: str = Query(...)) -> JSONResponse: |
| """ |
| Given tool documentation and a user question, generate a prompt |
| for the LLM to decide which tool to call and how. |
| Function Should be a valid query function name |
| """ |
|
|
| """ |
| Refer to tool_docs.json for full schema details. Create a valid Input for a query based off the given string and then directly make the query. |
| """ |
| with open(toolDocsPath, "r") as f: |
| doc_section = json.load(f) |
| tools_str = json.dumps(doc_section, indent=2) |
| |
|
|
| prompt = textwrap.dedent(f""" |
| You are a tool selector and payload generator. |
| Given the following tools and their documentation, choose the most appropriate tool to answer the user question. Then generate a valid JSON payload to call that tool. |
| |
| TOOLS: |
| {tools_str[function]} |
| |
| USER QUESTION: |
| {input} |
| |
| Respond ONLY in this JSON format: |
| {{ |
| "tool_name": "<tool_name>", |
| "payload": {{ |
| "columns": [...], |
| "filters": [...], |
| "metrics": [...], |
| "groupby": [...], |
| "orderby": [...], |
| "row_limit": <int> |
| }} |
| }} |
| """) |
| |
| return JSONResponse(content={"prompt": prompt.strip()}) |
| |
| @app.get("/products") |
| async def get_products(): |
| """ |
| Endpoint: Retrieve a list of products managed in Scribe: Should be called when asked for scribe products or a list of products. |
| |
| Args: Nothing |
| |
| Returns: |
| JSONResponse: Json list of products |
| """ |
| try: |
| return client.get_products() |
| except Exception as e: |
| logging.exception(f"Error in get_products endpoint {e}") |
|
|
| @app.get("/datasets") |
| async def get_datasets(): |
| """ |
| Endpoint: Retrieve all datasets: This should be called when anything about datasets is mentions. |
| |
| Args: Nothing |
| |
| Returns: |
| JSONResponse: Json list of datasets |
| """ |
| try: |
| return client.get_datasets() |
| except Exception as e: |
| logging.exception(f"Error in get_datasets endpoint: {e}") |
|
|
| @app.get("/{product_name}/vulns") |
| async def get_product_vulnerability(product_name: str): |
| """ |
| Endpoint: Retrieve vulnerabilities for a specific product: This should be called when a specifc product and the user is asking for vulnerabilities. |
| |
| Args: product name. |
| |
| Returns: |
| JSONResponse: Json list of product vulnerabilities |
| """ |
| products = get_products_list() |
| if product_name not in products: |
| logging.exception(f"Product name: {product_name} is not a valid product") |
| raise HTTPException(status_code=404, detail=str("Product not found")) |
| try: |
| return client.get_product_vulnerabilities(product_name) |
| except Exception as e: |
| logging.exception(f"Error in get_product_vulnerabilities endpoint: {e}") |
| |
|
|
| @app.get("/{product_name}/policy") |
| async def get_product_policy(product_name: str): |
| """ |
| Endpoint: Retrieve policy for a specific product: This should be called when a specifc product and the user is asking for policy. |
| |
| Args: product name. |
| |
| Returns: JSONResponse: Json list of product policy |
| """ |
| products = get_products_list() |
| if product_name not in products: |
| logging.exception(f"Product name: {product_name} is not a valid product") |
| raise HTTPException(status_code=404, detail=str("Product not found")) |
| try: |
| return client.get_policy_results(product_name) |
| except Exception as e: |
| logging.exception(f"Error in get_policy_results endpoint: {e}") |
|
|
| @app.post("/vulnerabilities") |
| async def query_vulnerabilities(input: dict): |
| """ |
| This tool is used to answer questions about vulnerabilities like: What are the vulnerabilities of my product? What are the vulnerabilities of a specific component? What are the vulnerabilities of a specific version of a component? Which products have a specific vulnerability, or a vulnerability with a specificy characteristic (such as has_key) |
| |
| Args: |
| { |
| "column": [], |
| "filters" [], |
| "metrics" [], |
| "orderby": [], |
| "row_limit": 20 |
| } |
| - "columns": list of columns to select. Use only these valid columns: |
| [('component_id', 'Dependency ID'), ('vulnerability_id', 'ID'), ('vector', 'Vector'), ('severity_string', 'Severity String'), ('base_severity_string', 'Base Severity String'), ('advisory_status', 'Advisory Status'), ('final_severity', 'Final Severity'), ('tool', 'Tool'), ('teamId', 'Teamid'), ('pipeline_run_id', 'Pipeline Run Id'), ('pkg', 'Pkg'), ('latest_package_version', 'Dependency Latest Version'), ('attestation_ids', 'Attestation IDs'), ('cvss_score', 'CVSS'), ('logical_app', 'Product'), ('vul_component_fixed_in_versions', 'Vulnerability Fix Version'), ('will_not_fix', 'Will Not Fix'), ('attestation_name', 'Attestation Name'), ('tag', 'Component Version'), ('targetName', 'Component'), ('base_score', 'Base Score'), ('source_layer', 'Layer'), ('text', 'Text'), ('is_latest_logical_version', 'Latest product version'), ('update', 'Update'), ('cwe_array', 'CWEs'), ('has_fix', 'Has Fix'), ('component_purl', 'Dependency PURL'), ('component_version', 'Dependency Version'), ('component_locations', 'Dependency Locations'), ('vul_date_changed', 'Vulnerability Last Change'), ('vul_component_created', 'Dependency Date'), ('epssProbability', 'EPSS'), ('vendor', 'Vendor'), ('direct_cwes', 'Direct Cwes'), ('advisory_modified', 'Advisory Modified'), ('hyper_links', 'Hyper Links'), ('advisory_text', 'Advisory Text'), ('cwes', 'CWE_Object'), ('advisory_justification', 'Advisory Justification'), ('component_dependency_type', 'Component Dependency Type'), ('cpes', 'Dependency CPEs'), ('has_kev', 'Has KEV'), ('version_timestamp', 'Version Timestamp'), ('severity', 'Original Severity'), ('epssPercentile', 'Epsspercentile'), ('epssScoreDate', 'Epssscoredate'), ('vul_adv_id', 'Vul Adv Id'), ('dependency_analysis', 'Dependency Analysis'), ('main_tool', 'Main Tool'), ('vul_published_on', 'Published'), ('vul_hyperlinks', 'References'), ('component_name', 'Dependency'), ('product_id', 'Product Id'), ('labels', 'Labels'), ('vul_source', 'Source'), ('targets', 'Targets'), ('logical_app_version', 'Version'), ('id', 'Id')] |
| |
| - "filters": list of filter objects with keys: |
| - "col": column name (must be valid) |
| - "op": operator: ==, !=, >, <, >=, <=, in, not in, is null, is not null |
| - "val": filter value |
| |
| - "metrics": list of metrics with keys: |
| - "label": output column name |
| - "expressionType": "SQL" or "SIMPLE" |
| - "sqlExpression": e.g. "COUNT(*)" |
| |
| - "groupby": list of valid column names |
| |
| - "orderby": list of [column_name, is_ascending_boolean] |
| |
| "row_limit": integer, max 30 |
| |
| Output ONLY the JSON object with no extra text. |
| """ |
| response = QueryResponse( |
| format="", |
| status=True, |
| status_info="Query executed successfully", |
| json_data={}, |
| str="" |
| ) |
| try: |
| result_str = client.query_vulnerabilities(json.dumps(input)) |
| logging.debug(result_str) |
| if result_str.startswith("Error"): |
| logging.exception("Error in query_vulnerabilities endpoint") |
| raise HTTPException(status_code=400, detail=result_str) |
| if result_str == f"\nNo data found": |
| response.status = False |
| response.status_info = "No data found" |
| elif format == "json": |
| response.json_data = result_str |
| else: |
| response.format = "str" |
| response.str = result_str |
| except HTTPException as he: |
| raise he |
| except Exception as e: |
| logging.exception(f"Error in query_vulnerabilities endpoint: {e}") |
| response.status = False |
| response.status_info = str(e) |
| return response |
| |
| @app.post("/products") |
| async def query_products(input: dict): |
| """ |
| To answer questions about products like: what are the/my products? What are the components of my product (the this case use the targetName column!)? What are the depdencieis of my product (search the component_name column)? Note that repositories and namespaces are part of the lineage table. |
| |
| Args: |
| { |
| "column": [], |
| "filters" [], |
| "metrics" [], |
| "orderby": [], |
| "row_limit": 20 |
| } |
| - "columns": list of columns to select. Use only these valid columns: |
| [('more', ' '), ('license_is_permissive_na', 'License Is Permissive Na'), ('teamId', 'Teamid'), ('component_version', 'Dep version'), ('license_expression', 'Licesne expression'), ('pkg', 'Package Manager'), ('rstatus', 'Version Status'), ('component_purl', 'PURL'), ('dependency_parents', 'Dependency Parents'), ('dependency_children', 'Dependency Children'), ('component_dependency_type', 'Component Dependency Type'), ('logical_app', 'Product'), ('logical_app_version', 'Version'), ('scorecard_id', 'Scorecard Id'), ('version_is_up_to_date', 'Up-to-date'), ('product_id', 'Product Id'), ('attestation_name', 'Attestation Name'), ('high_severity_cves', 'C+H Vulns'), ('maintenance_score', 'Maintenance Score'), ('component_name', 'Dependency'), ('labels', 'Labels'), ('licenses', 'Licenses'), ('attestation_ids', 'Attestation Ids'), ('component_external_ref', 'Component External Ref'), ('license_expression_na', 'License Expression Na'), ('component_publisher', 'Component Publisher'), ('component_description', 'Component Description'), ('targets_translated', 'SBOM Type Translated'), ('component_author', 'Component Author'), ('relation', 'Relation'), ('main_tool', 'Main Tool'), ('maintenance_score_int', 'Maintenance Score'), ('block', 'Dependency'), ('tools', 'Tools'), ('license_is_permissive', 'Is permissive'), ('scorecard_score', 'OpenSSF score'), ('targets', 'SBOM Type'), ('component_type', 'Component type'), ('source_layer', 'Layer'), ('component_locations', 'Component Locations'), ('component_id', 'Component Id'), ('id', 'Id'), ('base_layer', 'Base Image'), ('latest_version', 'Dep latest ver'), ('tag', 'Component version'), ('targetName', 'Component'), ('component_group', 'Component Group'), ('version_timestamp', 'Version Timestamp'), ('attestation_ids_cnt', 'Attestation Ids Cnt'), ('pipeline_run_id', 'Pipeline Run Id')] |
| - "filters": list of filter objects with keys: |
| - "col": column name (must be valid) |
| - "op": operator: ==, !=, >, <, >=, <=, in, not in, is null, is not null |
| - "val": filter value |
| |
| - "metrics": list of metrics with keys: |
| - "label": output column name |
| - "expressionType": "SQL" or "SIMPLE" |
| - "sqlExpression": e.g. "COUNT(*)" |
| |
| - "groupby": list of valid column names |
| |
| - "orderby": list of [column_name, is_ascending_boolean] |
| |
| "row_limit": integer, max 30 |
| |
| Output ONLY the JSON object with no extra text. |
| """ |
| response = QueryResponse( |
| format="", |
| status=True, |
| status_info="Query executed successfully", |
| json_data={}, |
| str="" |
| ) |
| try: |
| result_str = client.query_products(json.dumps(input)) |
| if result_str.startswith("Error"): |
| logging.exception("Error in query_products endpoint") |
| raise HTTPException(status_code=400, detail=result_str) |
| if result_str == f"\nNo data found": |
| response.status = False |
| response.status_info = "No data found" |
| elif format == "json": |
| response.json_data = result_str |
| else: |
| response.format = "str" |
| response.str = result_str |
| except HTTPException as he: |
| raise he |
| except Exception as e: |
| logging.exception(f"Error in query_products endpoint: {e}") |
| response.status = False |
| response.status_info = str(e) |
| return response |
|
|
| @app.post("/policy") |
| async def query_policy_results(input: dict): |
| response = QueryResponse( |
| format="", |
| status=True, |
| status_info="Query executed successfully", |
| json_data={}, |
| str="" |
| ) |
| try: |
| result_str = client.query_policy_results(json.dumps(input)) |
| if result_str.startswith("Error"): |
| logging.exception("Error in query_policy_results endpoint") |
| raise HTTPException(status_code=400, detail=result_str) |
| if result_str == f"\nNo data found": |
| response.status = False |
| response.status_info = "No data found" |
| elif format == "json": |
| response.json_data = result_str |
| else: |
| response.format = "str" |
| response.str = result_str |
| except HTTPException as he: |
| raise he |
| except Exception as e: |
| logging.exception(f"Error in query_policy_results endpoint: {e}") |
| response.status = False |
| response.status_info = str(e) |
| return response |
| |
| @app.post("/lineage") |
| async def query_lineage(input: dict): |
| """ |
| To answer software-factory related questions such as: What repos (or code repositories) are related/part of a product? or What are all the software-factory related components of a product (in this case query for asset_type and asset_name for the logical_app and logical_app_version)? |
| Args: |
| { |
| "column": [], |
| "filters" [], |
| "metrics" [], |
| "orderby": [], |
| "row_limit": 20 |
| } |
| - "columns": list of columns to select. Use only these valid columns: |
| [('platform_type', 'Platform Type'), ('properties', 'Info'), ('version_id', 'Version Id'), ('product_id', 'Product ID'), ('parent_name', 'Parent Name'), ('parent_type', 'Parent Type'), ('asset_name', 'Asset'), ('path', 'Path'), ('parent_id', 'Parent Id'), ('child_id', 'Child Id'), ('teamId', 'Teamid'), ('owner', 'Owner'), ('platform_name', 'Platform'), ('asset_type', 'Type'), ('external_id', 'External ID'), ('timestamp', 'Updated'), ('uri', 'URI'), ('prent_bloack', 'Parent'), ('more', ' More'), ('tooltip', 'Tooltip'), ('parent_external_id', 'Parent External ID'), ('logical_app', 'Product'), ('logical_app_version', 'Version'), ('asset_block', 'Asset'), ('g-source', 'G-Source'), ('g-target', 'G-Target')] - "filters": list of filter objects with keys: |
| - "col": column name (must be valid) |
| - "op": operator: ==, !=, >, <, >=, <=, in, not in, is null, is not null |
| - "val": filter value |
| |
| - "metrics": list of metrics with keys: |
| - "label": output column name |
| - "expressionType": "SQL" or "SIMPLE" |
| - "sqlExpression": e.g. "COUNT(*)" |
| |
| - "groupby": list of valid column names |
| |
| - "orderby": list of [column_name, is_ascending_boolean] |
| |
| "row_limit": integer, max 30 |
| |
| Output ONLY the JSON object with no extra text. |
| """ |
| response = QueryResponse( |
| format="", |
| status=True, |
| status_info="Query executed successfully", |
| json_data={}, |
| str="" |
| ) |
| try: |
| print(input) |
| result_str = client.query_lineage(json.dumps(input)) |
| if result_str.startswith("Error"): |
| logging.exception("Error in query_lineage endpoint") |
| raise HTTPException(status_code=400, detail=result_str) |
| if result_str == f"\nNo data found": |
| response.status = False |
| response.status_info = "No data found" |
| elif format == "json": |
| response.json_data = result_str |
| else: |
| response.format = "str" |
| response.str = result_str |
| except HTTPException as he: |
| raise he |
| except Exception as e: |
| logging.exception(f"Error in query_lineage endpoint: {e}") |
| response.status = False |
| response.status_info = str(e) |
| return response |
|
|
| @app.post("/risk") |
| async def query_risk(input: dict): |
| response = QueryResponse( |
| format="", |
| status=True, |
| status_info="Query executed successfully", |
| json_data={}, |
| str="" |
| ) |
| try: |
| print(input) |
| result_str = client.query_risk(json.dumps(input)) |
| if result_str.startswith("Error"): |
| logging.exception("Error in query_risk endpoint") |
| raise HTTPException(status_code=400, detail=result_str) |
| if result_str == f"\nNo data found": |
| response.status = False |
| response.status_info = "No data found" |
| elif format == "json": |
| response.json_data = result_str |
| else: |
| response.format = "str" |
| response.str = result_str |
| except HTTPException as he: |
| raise he |
| except Exception as e: |
| logging.exception(f"Error in query_risk endpoint: {e}") |
| response.status = False |
| response.status_info = str(e) |
| return response |
|
|
| @app.post("/findings") |
| async def query_findings(input: dict): |
| response = QueryResponse( |
| format="", |
| status=True, |
| status_info="Query executed successfully", |
| json_data={}, |
| str="" |
| ) |
| try: |
| print(input) |
| result_str = client.query_findings(json.dumps(input)) |
| if result_str.startswith("Error"): |
| logging.exception("Error in query_findings endpoint") |
| raise HTTPException(status_code=400, detail=result_str) |
| if result_str == f"\nNo data found": |
| response.status = False |
| response.status_info = "No data found" |
| elif format == "json": |
| response.json_data = result_str |
| else: |
| response.format = "str" |
| response.str = result_str |
| except HTTPException as he: |
| raise he |
| except Exception as e: |
| logging.exception(f"Error in query_findings endpoint: {e}") |
| response.status = False |
| response.status_info = str(e) |
| return response |
