diff --git a/examples/tools/cve_check.py b/examples/tools/cve_check.py index f9aa14c..cccf7ce 100644 --- a/examples/tools/cve_check.py +++ b/examples/tools/cve_check.py @@ -6,7 +6,9 @@ if __name__ == "__main__": ipf = IPFClient() # ipf = IPFClient('https://demo3.ipfabric.io/', token='token', verify=False, timeout=15) - vuln = Vulnerabilities(ipf) + vuln = Vulnerabilities(ipf, timeout=30, cve_limit=20) + # Increasing cve_limit will increase the amount of time NIST responds + device = vuln.check_device('L47R6') pprint(device) diff --git a/ipfabric/tools/nist.py b/ipfabric/tools/nist.py index 2c826a1..9993bed 100644 --- a/ipfabric/tools/nist.py +++ b/ipfabric/tools/nist.py @@ -9,12 +9,16 @@ class CVEs(BaseModel): class NIST(Client): + def __init__(self, timeout, cve_limit): + super().__init__(base_url='https://services.nvd.nist.gov/rest/json/cves/1.0', timeout=timeout) + self.cve_limit = cve_limit + @property def params(self): return { 'cpeMatchString': 'cpe:2.3:*:', 'startIndex': 0, - 'resultsPerPage': 50 + 'resultsPerPage': self.cve_limit } def check_cve(self, vendor: str, family: str, version: str): diff --git a/ipfabric/tools/vulnerabilities.py b/ipfabric/tools/vulnerabilities.py index 8e4ec1d..ca928f9 100644 --- a/ipfabric/tools/vulnerabilities.py +++ b/ipfabric/tools/vulnerabilities.py @@ -15,9 +15,9 @@ class Version(BaseModel): class Vulnerabilities: - def __init__(self, ipf): + def __init__(self, ipf, timeout: int = 30, cve_limit: int = 20): self.ipf = ipf - self.nist = NIST(base_url='https://services.nvd.nist.gov/rest/json/cves/1.0', timeout=30) + self.nist = NIST(timeout=timeout, cve_limit=cve_limit) def __del__(self): try: