diff --git a/greenbone/scap/cpe/cli/download.py b/greenbone/scap/cpe/cli/download.py index 46995fd..948e885 100644 --- a/greenbone/scap/cpe/cli/download.py +++ b/greenbone/scap/cpe/cli/download.py @@ -424,10 +424,14 @@ async def download(console: Console, error_console: Console) -> None: ) if run_time_file: + if until: + run_time = until + else: + run_time = datetime.now() # ensure directories exist run_time_file.parent.mkdir(parents=True, exist_ok=True) run_time_file.write_text( - f"{until.isoformat()}\n", + f"{run_time.isoformat()}\n", encoding="utf8", # type: ignore ) console.log(f"Wrote run time to {run_time_file.absolute()}.") diff --git a/greenbone/scap/cpe_match/db/manager.py b/greenbone/scap/cpe_match/db/manager.py index dc9cea5..d7b7196 100644 --- a/greenbone/scap/cpe_match/db/manager.py +++ b/greenbone/scap/cpe_match/db/manager.py @@ -107,7 +107,7 @@ async def add_cpe_match_strings( ) await self._insert_foreign_data(transaction, match_strings) - self._cpes = [] + self._cpe_match_strings = [] async def _insert_foreign_data( self, @@ -150,8 +150,8 @@ async def find( index: int | None = None, last_modification_start_date: datetime | None = None, last_modification_end_date: datetime | None = None, - published_start_date: datetime | None = None, - published_end_date: datetime | None = None, + created_start_date: datetime | None = None, + created_end_date: datetime | None = None, ) -> AsyncIterator[CPEMatchStringDatabaseModel]: clauses = [] @@ -171,13 +171,13 @@ async def find( CPEMatchStringDatabaseModel.last_modified <= last_modification_end_date ) - if published_start_date: + if created_start_date: clauses.append( - CPEMatchStringDatabaseModel.published >= published_start_date + CPEMatchStringDatabaseModel.created >= created_start_date ) - if published_end_date: + if created_end_date: clauses.append( - CPEMatchStringDatabaseModel.published <= published_end_date + CPEMatchStringDatabaseModel.created <= created_end_date ) statement = ( @@ -221,8 +221,8 @@ async def count( match_criteria_id: str | None, last_modification_start_date: datetime | None = None, last_modification_end_date: datetime | None = None, - published_start_date: datetime | None = None, - published_end_date: datetime | None = None, + created_start_date: datetime | None = None, + created_end_date: datetime | None = None, ) -> int: clauses = [] @@ -242,13 +242,13 @@ async def count( CPEMatchStringDatabaseModel.last_modified <= last_modification_end_date ) - if published_start_date: + if created_start_date: clauses.append( - CPEMatchStringDatabaseModel.published >= published_start_date + CPEMatchStringDatabaseModel.created >= created_start_date ) - if published_end_date: + if created_end_date: clauses.append( - CPEMatchStringDatabaseModel.published <= published_end_date + CPEMatchStringDatabaseModel.created <= created_end_date ) statement = select( diff --git a/greenbone/scap/cpe_match/json.py b/greenbone/scap/cpe_match/json.py index 273b39c..bee0b85 100644 --- a/greenbone/scap/cpe_match/json.py +++ b/greenbone/scap/cpe_match/json.py @@ -7,6 +7,7 @@ from dataclasses import asdict, dataclass from datetime import datetime from pathlib import Path +from typing import Sequence from pontos.nvd.models.cpe_match_string import CPEMatchString from rich.console import Console @@ -65,7 +66,7 @@ def __init__( storage_path: Path, *, compress: bool = False, - schema_path: Path = None, + schema_path: Path | None = None, raise_error_on_validation=False, ): super().__init__( @@ -88,7 +89,9 @@ def add_match_string(self, match_string: CPEMatchString) -> None: MatchStringItem(match_string=match_string) ) - def add_match_strings(self, match_strings: list[CPEMatchString]) -> None: + def add_match_strings( + self, match_strings: Sequence[CPEMatchString] + ) -> None: for match_string in match_strings: self.add_match_string(match_string) diff --git a/greenbone/scap/cpe_match/producer/nvd_api.py b/greenbone/scap/cpe_match/producer/nvd_api.py index 7199f69..46960b8 100644 --- a/greenbone/scap/cpe_match/producer/nvd_api.py +++ b/greenbone/scap/cpe_match/producer/nvd_api.py @@ -9,6 +9,8 @@ from rich.console import Console from rich.progress import Progress +from ...cli import DEFAULT_RETRIES +from ...errors import ScapError from ...generic_cli.producer.nvd_api import NvdApiProducer from ..cli.processor import CPE_MATCH_TYPE_PLURAL @@ -69,12 +71,12 @@ def __init__( error_console: Console, progress: Progress, *, + retry_attempts: int = DEFAULT_RETRIES, nvd_api_key: str | None = None, - retry_attempts: int = None, - request_results: int = None, + request_results: int | None = None, request_filter_opts: dict = {}, start_index: int = 0, - verbose: int = None, + verbose: int | None = None, ): """ Constructor for a CPE match string NVD API producer. @@ -90,19 +92,21 @@ def __init__( start_index: index/offset of the first item to request verbose: Verbosity level of log messages. """ + self._nvd_api: CPEMatchApi + super().__init__( console, error_console, progress, - nvd_api_key=nvd_api_key, retry_attempts=retry_attempts, + nvd_api_key=nvd_api_key, request_results=request_results, request_filter_opts=request_filter_opts, start_index=start_index, verbose=verbose, ) - def _create_nvd_api(self, nvd_api_key: str) -> CPEMatchApi: + def _create_nvd_api(self, nvd_api_key: str | None) -> CPEMatchApi: """ Callback used by the constructor to create the NVD API object that can be queried for CPE match strings. @@ -123,6 +127,9 @@ async def _create_nvd_results(self) -> NVDResults[CPEMatchString]: Returns: The new `NVDResults` object. """ + if self._queue is None: + raise ScapError("No queue has been assigned") + return await self._nvd_api.cpe_matches( last_modified_start_date=self._request_filter_opts.get( "last_modified_start_date" diff --git a/greenbone/scap/cpe_match/worker/db.py b/greenbone/scap/cpe_match/worker/db.py index c8ab066..0b8d9e2 100644 --- a/greenbone/scap/cpe_match/worker/db.py +++ b/greenbone/scap/cpe_match/worker/db.py @@ -92,6 +92,8 @@ def __init__( echo_sql: Whether to print SQL statements. verbose: Verbosity level of log messages. """ + self._manager: CPEMatchStringDatabaseManager + super().__init__( console, error_console, diff --git a/greenbone/scap/cpe_match/worker/json.py b/greenbone/scap/cpe_match/worker/json.py index 06bd8de..5b7c309 100644 --- a/greenbone/scap/cpe_match/worker/json.py +++ b/greenbone/scap/cpe_match/worker/json.py @@ -111,11 +111,11 @@ async def _handle_chunk(self, chunk: Sequence[CPEMatchString]): """ self._json_manager.add_match_strings(chunk) - async def loop_end(self) -> None: + async def _loop_end(self) -> None: """ Callback handling the exiting the main worker loop. Makes the JSON manager write the document to the file. """ self._json_manager.write() - await super().loop_end() + await super()._loop_end() diff --git a/greenbone/scap/cve/cli/download.py b/greenbone/scap/cve/cli/download.py index c506355..16b04e2 100644 --- a/greenbone/scap/cve/cli/download.py +++ b/greenbone/scap/cve/cli/download.py @@ -432,10 +432,14 @@ async def download(console: Console, error_console: Console): ) if run_time_file: + if until: + run_time = until + else: + run_time = datetime.now() # ensure directories exist run_time_file.parent.mkdir(parents=True, exist_ok=True) run_time_file.write_text( - f"{until.isoformat()}\n", + f"{run_time.isoformat()}\n", encoding="utf8", # type: ignore ) console.log(f"Wrote run time to {run_time_file.absolute()}.") diff --git a/greenbone/scap/data_utils/json.py b/greenbone/scap/data_utils/json.py index 9897d0e..1f6e21b 100644 --- a/greenbone/scap/data_utils/json.py +++ b/greenbone/scap/data_utils/json.py @@ -123,7 +123,7 @@ class JsonManager: def __init__( self, error_console: Console, - schema_path: Path = None, + schema_path: Path | None = None, raise_error_on_validation=False, ): """ diff --git a/greenbone/scap/generic_cli/processor.py b/greenbone/scap/generic_cli/processor.py index 7b1497d..9e65f39 100644 --- a/greenbone/scap/generic_cli/processor.py +++ b/greenbone/scap/generic_cli/processor.py @@ -118,7 +118,7 @@ def __init__( self._worker.set_queue(self._queue) self._verbose: int = ( - verbose if not None else self._arg_defaults["verbose"], + verbose if verbose is not None else self._arg_defaults["verbose"] ) "Verbosity level of log messages." diff --git a/greenbone/scap/generic_cli/producer/base.py b/greenbone/scap/generic_cli/producer/base.py index 9e01a52..14a2c49 100644 --- a/greenbone/scap/generic_cli/producer/base.py +++ b/greenbone/scap/generic_cli/producer/base.py @@ -4,7 +4,7 @@ from abc import ABC, abstractmethod from argparse import ArgumentParser -from typing import AsyncContextManager, Generic, TypeVar +from typing import Any, AsyncContextManager, Generic, TypeVar from rich.console import Console from rich.progress import Progress @@ -30,10 +30,10 @@ class BaseScapProducer(Generic[T], AsyncContextManager, ABC): e.g. `BaseScapProducer[CPE]` will be a producer handling CPE objects. """ - _item_type_plural = "SCAP items" + _item_type_plural: str = "SCAP items" "Plural form of the type of items to use in log messages" - _arg_defaults = { + _arg_defaults: dict[str, Any] = { "verbose": DEFAULT_VERBOSITY, } "Default values for optional arguments." @@ -79,7 +79,7 @@ def __init__( self._verbose = verbose if not None else self._arg_defaults["verbose"] "Verbosity level of log messages." - self._queue: ScapChunkQueue[T] | None = None + self._queue: ScapChunkQueue[T] "Queue chunks of SCAP items are added to." @abstractmethod @@ -110,7 +110,7 @@ async def run_loop( It should also create a task for the `progress` object and update it regularly. """ - self._queue.set_producer_finished() + pass def set_queue(self, queue: ScapChunkQueue[T]) -> None: """ diff --git a/greenbone/scap/generic_cli/producer/nvd_api.py b/greenbone/scap/generic_cli/producer/nvd_api.py index 4f08e36..7777a80 100644 --- a/greenbone/scap/generic_cli/producer/nvd_api.py +++ b/greenbone/scap/generic_cli/producer/nvd_api.py @@ -5,7 +5,7 @@ from argparse import ArgumentParser, Namespace from datetime import datetime from pathlib import Path -from typing import Generic, TypeVar +from typing import Any, Generic, TypeVar import httpx import stamina @@ -98,7 +98,9 @@ def add_args_to_parser( ) @staticmethod - def since_from_args(args: Namespace, error_console: Console) -> datetime: + def since_from_args( + args: Namespace, error_console: Console + ) -> datetime | None: """ Gets the lower limit for the modification time from the given command line arguments, reading the time from a file if the @@ -133,8 +135,8 @@ def __init__( error_console: Console, progress: Progress, *, + retry_attempts: int = DEFAULT_RETRIES, nvd_api_key: str | None = None, - retry_attempts: int, request_results: int | None = None, request_filter_opts: dict = {}, start_index: int = 0, @@ -167,26 +169,26 @@ def __init__( self._additional_retry_attempts: int = retry_attempts "Number of retries after fetching initial data." - self._request_results: int = request_results + self._request_results: int | None = request_results "Maximum number of results to request from the API." - self._request_filter_opts: dict = request_filter_opts + self._request_filter_opts: dict[str, Any] = request_filter_opts "Filter options to pass to the API requests." self._start_index: int = start_index "Index/offset of the first item to request." - self._nvd_api_key = nvd_api_key + self._nvd_api_key: str | None = nvd_api_key "API key to use for the requests to allow faster requests." - self._nvd_api = self._create_nvd_api(nvd_api_key) + self._nvd_api: NVDApi = self._create_nvd_api(nvd_api_key) "The NVD API object used for querying SCAP items." - self._results = None + self._results: NVDResults[T] "The NVD results object created by the API to get the SCAP items from." @abstractmethod - def _create_nvd_api(self, nvd_api_key: str) -> NVDApi: + def _create_nvd_api(self, nvd_api_key: str | None) -> NVDApi: """ Callback used by the constructor to create the NVD API object that can be queried for SCAP items. @@ -199,7 +201,7 @@ def _create_nvd_api(self, nvd_api_key: str) -> NVDApi: pass @abstractmethod - def _create_nvd_results(self) -> NVDResults[T]: + async def _create_nvd_results(self) -> NVDResults[T]: """ Callback used during `fetch_initial_data` to get the `NVDResults` object the SCAP items will be fetched from. diff --git a/greenbone/scap/generic_cli/worker/base.py b/greenbone/scap/generic_cli/worker/base.py index 11470ee..32e11fb 100644 --- a/greenbone/scap/generic_cli/worker/base.py +++ b/greenbone/scap/generic_cli/worker/base.py @@ -5,7 +5,7 @@ import asyncio from abc import ABC, abstractmethod from argparse import ArgumentParser -from typing import AsyncContextManager, Generic, Sequence, TypeVar +from typing import Any, AsyncContextManager, Generic, Sequence, TypeVar from rich.console import Console from rich.progress import Progress, TaskID @@ -27,10 +27,10 @@ class BaseScapWorker(Generic[T], AsyncContextManager, ABC): e.g. `BaseScapWorker[CPE]` will be a producer handling CPE objects. """ - _item_type_plural = "SCAP items" + _item_type_plural: str = "SCAP items" "Plural form of the type of items to use in log messages." - _arg_defaults = { + _arg_defaults: dict[str, Any] = { "verbose": DEFAULT_VERBOSITY, } "Default values for optional arguments." @@ -77,7 +77,7 @@ def __init__( self._verbose = verbose if not None else self._arg_defaults["verbose"] "Verbosity level of log messages." - self._queue: ScapChunkQueue[T] | None = None + self._queue: ScapChunkQueue[T] "Queue the worker will get chunks of SCAP items from." self._progress_task: TaskID | None = None @@ -106,7 +106,7 @@ async def _loop_start(self) -> None: total=self._queue.total_items, ) - async def loop_step_end(self) -> None: + async def _loop_step_end(self) -> None: """ Callback handling the end of one iteration of the main worker loop. """ @@ -116,7 +116,7 @@ async def loop_step_end(self) -> None: f"{self._item_type_plural}" ) - async def loop_end(self) -> None: + async def _loop_end(self) -> None: """ Callback handling the exiting the main worker loop. """ @@ -134,6 +134,9 @@ async def run_loop(self) -> None: It will also call `loop_step_end` after each iteration of the loop and `loop_end` after exiting the loop. """ + if self._queue is None: + raise ScapError("No queue has been assigned") + await self._loop_start() while self._queue.more_chunks_expected(): try: @@ -155,9 +158,9 @@ async def run_loop(self) -> None: self._queue.chunk_processed() - await self.loop_step_end() + await self._loop_step_end() - await self.loop_end() + await self._loop_end() def set_queue(self, queue: ScapChunkQueue[T]) -> None: """ diff --git a/greenbone/scap/generic_cli/worker/db.py b/greenbone/scap/generic_cli/worker/db.py index 168dc99..16c7f01 100644 --- a/greenbone/scap/generic_cli/worker/db.py +++ b/greenbone/scap/generic_cli/worker/db.py @@ -5,7 +5,7 @@ import os from abc import abstractmethod from argparse import ArgumentParser -from typing import AsyncContextManager, Sequence, TypeVar +from typing import AsyncContextManager, Sequence, Type, TypeVar from rich.console import Console from rich.progress import Progress @@ -59,7 +59,7 @@ class ScapDatabaseWriteWorker(BaseScapWorker[T]): @classmethod def add_args_to_parser( - cls: type, + cls: Type["ScapDatabaseWriteWorker"], parser: ArgumentParser, ): """ @@ -168,7 +168,8 @@ def __init__( or self._arg_defaults["database_host"] ) try: - env_database_port = int(os.environ.get("DATABASE_PORT")) + port_str = os.environ.get("DATABASE_PORT") + env_database_port = int(port_str) if port_str else None except TypeError: env_database_port = None database_port = ( @@ -182,11 +183,13 @@ def __init__( ) if not database_user: - raise CLIError(f"Missing user for {self.item_type_plural} database") + raise CLIError( + f"Missing user for {self._item_type_plural} database" + ) if not database_password: raise CLIError( - f"Missing password for {self.item_type_plural} database" + f"Missing password for {self._item_type_plural} database" ) self._database = PostgresDatabase( diff --git a/greenbone/scap/generic_cli/worker/json.py b/greenbone/scap/generic_cli/worker/json.py index 39ef7ab..9f48d4e 100644 --- a/greenbone/scap/generic_cli/worker/json.py +++ b/greenbone/scap/generic_cli/worker/json.py @@ -5,7 +5,7 @@ from abc import ABC from argparse import ArgumentParser from pathlib import Path -from typing import TypeVar +from typing import Type, TypeVar from rich.console import Console from rich.progress import Progress @@ -38,7 +38,7 @@ class ScapJsonWriteWorker(BaseScapWorker[T], ABC): @classmethod def add_args_to_parser( - cls: type, + cls: Type["ScapJsonWriteWorker"], parser: ArgumentParser, ): """