Spaces:
Sleeping
Sleeping
File size: 5,960 Bytes
4531c67 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
import requests
from langchain.tools import tool
import json
from datetime import datetime, timedelta
def get_current_formatted_date():
# Get the current date and time
now = datetime.now()
# Format the date and time in the desired format
formatted_date = now.strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]
return formatted_date
def get_yesterday_formatted_date():
# Get the current date and time
now = datetime.now()
# Calculate yesterday's date and time
yesterday = now - timedelta(days=2)
# Format the date and time in the desired format
formatted_date = yesterday.strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]
return formatted_date
def format_cve(cve_response):
"""
Takes a dictionary representing the API response and formats each CVE entry into a structured prompt for LLM.
Parameters:
- cve_response: A dictionary representing the API response containing CVE information.
Returns:
- A list of strings, each a formatted prompt for LLM based on the CVE entries in the response.
"""
formatted_prompts = "Use these CVE Search Results to answer user question:\n\n"
if len(cve_response['vulnerabilities']) == 0:
formatted_prompts = "No CVEs found matching the search criteria."
return formatted_prompts
for vulnerability in cve_response['vulnerabilities']:
cve = vulnerability.get('cve', {})
# prompt = f"Explain {cve.get('id', 'N/A')} in simple terms:\n\n"
prompt = f"- CVE ID: {cve.get('id', 'N/A')}\n"
prompt += f"- Status: {cve.get('vulnStatus', 'Unknown')}\n"
descriptions = cve.get('descriptions', [])
description_text = descriptions[0].get('value', 'No description available.') if descriptions else "No description available."
prompt += f"- Description: {description_text}\n"
if 'metrics' in cve and 'cvssMetricV2' in cve['metrics']:
cvss_metrics = cve['metrics']['cvssMetricV2'][0]
prompt += f"- CVSS Score: {cvss_metrics.get('cvssData', {}).get('baseScore', 'Not available')} ({cvss_metrics.get('baseSeverity', 'Unknown')})\n"
else:
prompt += "- CVSS Score: Not available\n"
configurations = cve.get('configurations', {})
for conf in configurations:
nodes = conf.get('nodes', [])
affected_configs = []
for node in nodes:
for cpe_match in node.get('cpeMatch', []):
if cpe_match.get('vulnerable', False):
affected_configs.append(cpe_match.get('criteria', 'Not specified'))
prompt += f"- Affected Configurations: {', '.join(affected_configs) if affected_configs else 'Not specified'}\n"
references = cve.get('references', [])
ref_urls = ', '.join([ref.get('url', 'No URL') for ref in references])
prompt += f"- References: {ref_urls if references else 'No references available.'}\n"
formatted_prompts += prompt+"\n\n"
print(formatted_prompts)
# formatted_prompts += "\nSummarize the vulnerability, its impact, and any known mitigation strategies."
return formatted_prompts
class CVESearchTool():
@tool("CVE search Tool")
def cvesearch(keyword: str, date: str = None):
"""
Searches for CVEs based on a keyword or phrase and returns the results in JSON format.
Parameters:
- keyword (str): A word or phrase to search for in the CVE descriptions.
- date (str): An optional date to include in the search query.
Returns:
- JSON: A list of CVEs matching the keyword search.
"""
if date:
keyword = f"{keyword} {date}"
# Encode the spaces in the keyword(s) as "%20" for the URL
keyword_encoded = keyword.replace(" ", "%20")
# keyword_encoded = keyword_encoded.join(" 2023")
# Construct the URL for the API request
url = f"https://services.nvd.nist.gov/rest/json/cves/2.0?keywordSearch={keyword_encoded}&resultsPerPage=5"
try:
# Send the request to the NVD API
response = requests.get(url)
# Check if the request was successful
if response.status_code == 200:
# Return the JSON response
return format_cve(response.json())
else:
return {"error": "Failed to fetch data from the NVD API.", "status_code": response.status_code}
except Exception as e:
return {"error": str(e)}
@tool("Fetch last 5 CVEs")
def fetchcve(keyword: str):
"""
Fetches the last 5 CVEs based on a keyword or phrase and returns the results in JSON format.
Use this exclusively to get the latest CVEs, or for todays CVEs.
Parameters:
- keyword (str): A word or phrase to search for in the CVE descriptions.
- date (str): An optional date to include in the search query.
Returns:
- JSON: A list of CVEs matching the keyword search.
"""
# Encode the spaces in the keyword(s) as "%20" for the URL
keyword_encoded = keyword.replace(" ", "%20")
# keyword_encoded = keyword_encoded.join(" 2023")
# Construct the URL for the API request
url = f"https://services.nvd.nist.gov/rest/json/cves/2.0/?pubStartDate={get_yesterday_formatted_date()}&pubEndDate={get_current_formatted_date()}&keywordSearch={keyword_encoded}&resultsPerPage=3"
print(url)
try:
# Send the request to the NVD API
response = requests.get(url)
# Check if the request was successful
if response.status_code == 200:
# Return the JSON response
return format_cve(response.json())
else:
return {"error": "Failed to fetch data from the NVD API.", "status_code": response.status_code}
except Exception as e:
return {"error": str(e)} |