diff --git a/pyvulnerabilitylookup/api.py b/pyvulnerabilitylookup/api.py index 2c59088..7fe3274 100644 --- a/pyvulnerabilitylookup/api.py +++ b/pyvulnerabilitylookup/api.py @@ -4,7 +4,7 @@ from importlib.metadata import version from pathlib import PurePosixPath -from typing import Any +from typing import Any, Dict from urllib.parse import urljoin, urlparse import requests @@ -12,7 +12,7 @@ class PyVulnerabilityLookup(): - def __init__(self, root_url: str, useragent: str | None=None, + def __init__(self, root_url: str, useragent: str | None=None, token: str | None=None, *, proxies: dict[str, str] | None=None) -> None: '''Query a specific instance. @@ -28,6 +28,9 @@ def __init__(self, root_url: str, useragent: str | None=None, self.root_url += '/' self.session = requests.session() self.session.headers['user-agent'] = useragent if useragent else f'PyProject / {version("pyvulnerabilitylookup")}' + self.session.headers['X-API-KEY'] = token if token else '' + self.session.headers['Accept'] = 'application/json' + self.session.headers['Content-Type'] = 'application/json' if proxies: self.session.proxies.update(proxies) @@ -96,11 +99,20 @@ def get_vendor_product_vulnerabilities(self, vendor: str, product: str) -> list[ # NOTE: endpoints /api/cve/*, /api/dbInfo, /api/last are alises for backward compat. + def create_comment(self, comment: Dict[str, Any]) -> Dict[str, Any]: + '''Create a comment. + + :param comment: The comment + ''' + r = self.session.post(urljoin(self.root_url, str(PurePosixPath('api', 'comment'))), + json=comment) + return r.json() + def get_comments(self, uuid: str | None = None, vuln_id: str | None = None, author: str | None = None) -> dict[str, Any]: '''Get comment(s) - :param uuid: The UUID a specific comment + :param uuid: The UUID of a specific comment :param vuln_id: The vulnerability ID to get comments of :param author: The author of the comment(s) ''' @@ -108,6 +120,23 @@ def get_comments(self, uuid: str | None = None, vuln_id: str | None = None, params={'uuid': uuid, 'vuln_id': vuln_id, 'author': author}) return r.json() + def delete_comment(self, comment_uuid: str) -> int: + '''Delete a comment. + + :param comment_uuid: The comment UUID + ''' + r = self.session.delete(urljoin(self.root_url, str(PurePosixPath('api', 'comment', comment_uuid)))) + return r.status_code + + def create_bundle(self, bundle: Dict[str, Any]) -> Dict[str, Any]: + '''Create a bundle. + + :param bundle: The bundle + ''' + r = self.session.post(urljoin(self.root_url, str(PurePosixPath('api', 'bundle'))), + json=bundle) + return r.json() + def get_bundles(self, uuid: str | None = None, vuln_id: str | None = None, author: str | None = None) -> dict[str, Any]: '''Get bundle(s) @@ -119,3 +148,11 @@ def get_bundles(self, uuid: str | None = None, vuln_id: str | None = None, r = self.session.get(urljoin(self.root_url, str(PurePosixPath('api', 'bundle'))), params={'uuid': uuid, 'vuln_id': vuln_id, 'author': author}) return r.json() + + def delete_bundle(self, bundle_uuid: str) -> int: + '''Delete a bundle. + + :param bundle_uuid: The bundle UUID + ''' + r = self.session.delete(urljoin(self.root_url, str(PurePosixPath('api', 'bundle', bundle_uuid)))) + return r.status_code diff --git a/tests/test_web.py b/tests/test_web.py index 73d8098..c41a3eb 100644 --- a/tests/test_web.py +++ b/tests/test_web.py @@ -2,6 +2,7 @@ import unittest import time +import os from pyvulnerabilitylookup import PyVulnerabilityLookup @@ -9,7 +10,8 @@ class TestPublic(unittest.TestCase): def setUp(self) -> None: - self.client = PyVulnerabilityLookup(root_url="https://vulnerability.circl.lu") + token = os.getenv("API_KEY", "") + self.client = PyVulnerabilityLookup(root_url="https://vulnerability.circl.lu", token=token) # Test default