From 53be3c224ab2c686b32141d7cfdfc49408c4ef6e Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 24 Jan 2024 16:14:21 -0500 Subject: [PATCH 001/139] Add cyhy_db package and models Delete example package and files Update Python version to 3.12 in build workflows Update package name and description in setup.py Update dependencies in setup.py --- .coveragerc | 2 +- .github/workflows/build.yml | 7 +- setup.py | 25 +++--- src/cyhy_db/_version.py | 1 + src/cyhy_db/models/cve.py | 43 ++++++++++ src/cyhy_db/utils/__init__.py | 0 src/example/__init__.py | 10 --- src/example/__main__.py | 5 -- src/example/_version.py | 3 - src/example/data/secret.txt | 1 - src/example/example.py | 103 ------------------------ tests/conftest.py | 92 ++++++++++++++++++++++ tests/test_cve.py | 35 +++++++++ tests/test_example.py | 144 ---------------------------------- 14 files changed, 189 insertions(+), 282 deletions(-) create mode 100644 src/cyhy_db/_version.py create mode 100644 src/cyhy_db/models/cve.py create mode 100644 src/cyhy_db/utils/__init__.py delete mode 100644 src/example/__init__.py delete mode 100644 src/example/__main__.py delete mode 100644 src/example/_version.py delete mode 100644 src/example/data/secret.txt delete mode 100644 src/example/example.py create mode 100644 tests/test_cve.py delete mode 100644 tests/test_example.py diff --git a/.coveragerc b/.coveragerc index d315b87..de936df 100644 --- a/.coveragerc +++ b/.coveragerc @@ -2,7 +2,7 @@ # https://coverage.readthedocs.io/en/latest/config.html [run] -source = src/example +source = src/cyhy_db omit = branch = true diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e429274..3a7ef12 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -62,7 +62,7 @@ jobs: - id: setup-python uses: actions/setup-python@v5 with: - python-version: ${{ steps.setup-env.outputs.python-version }} + python-version: "3.12" # We need the Go version and Go cache location for the actions/cache step, # so the Go installation must happen before that. - id: setup-go @@ -180,7 +180,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.7" - "3.8" - "3.9" - "3.10" @@ -248,7 +247,7 @@ jobs: - id: setup-python uses: actions/setup-python@v5 with: - python-version: ${{ steps.setup-env.outputs.python-version }} + python-version: "3.12" - uses: actions/cache@v3 env: BASE_CACHE_KEY: "${{ github.job }}-${{ runner.os }}-\ @@ -286,7 +285,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.7" - "3.8" - "3.9" - "3.10" @@ -341,7 +339,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.7" - "3.8" - "3.9" - "3.10" diff --git a/setup.py b/setup.py index fdb21eb..00149ef 100644 --- a/setup.py +++ b/setup.py @@ -42,10 +42,10 @@ def get_version(version_file): setup( - name="example", + name="cyhy-db", # Versions should comply with PEP440 - version=get_version("src/example/_version.py"), - description="Example Python library", + version=get_version("src/cyhy_db/_version.py"), + description="CyHy Database Python library", long_description=readme(), long_description_content_type="text/markdown", # Landing page for CISA's cybersecurity mission @@ -75,7 +75,6 @@ def get_version(version_file): # that you indicate whether you support Python 2, Python 3 or both. "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", @@ -83,17 +82,23 @@ def get_version(version_file): "Programming Language :: Python :: 3.12", "Programming Language :: Python :: Implementation :: CPython", ], - python_requires=">=3.7", + python_requires=">=3.8", # What does your project relate to? - keywords="skeleton", + keywords=["cyhy", "database"], packages=find_packages(where="src"), package_dir={"": "src"}, - package_data={"example": ["data/*.txt"]}, + package_data={}, py_modules=[splitext(basename(path))[0] for path in glob("src/*.py")], include_package_data=True, - install_requires=["docopt", "schema", "setuptools >= 24.2.0"], + install_requires=[ + "docopt == 0.6.2", + "odmantic == 1.0.0", + "schema == 0.7.5", + "setuptools >= 69.0.3", + ], extras_require={ "test": [ + "pytest-asyncio", "coverage", # coveralls 1.11.0 added a service number for calls from # GitHub Actions. This caused a regression which resulted in a 422 @@ -102,11 +107,11 @@ def get_version(version_file): # 1.11.1 fixed this issue, but to ensure expected behavior we'll pin # to never grab the regression version. "coveralls != 1.11.0", + "docker == 7.0.0", "pre-commit", "pytest-cov", "pytest", ] }, - # Conveniently allows one to run the CLI tool as `example` - entry_points={"console_scripts": ["example = example.example:main"]}, + entry_points={}, ) diff --git a/src/cyhy_db/_version.py b/src/cyhy_db/_version.py new file mode 100644 index 0000000..b3c06d4 --- /dev/null +++ b/src/cyhy_db/_version.py @@ -0,0 +1 @@ +__version__ = "0.0.1" \ No newline at end of file diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py new file mode 100644 index 0000000..f8dfcb8 --- /dev/null +++ b/src/cyhy_db/models/cve.py @@ -0,0 +1,43 @@ +from odmantic import Model, Field +from pydantic import ValidationInfo, field_validator +from typing import Literal + +class CVE(Model): + id: str = Field(primary_field=True) + cvss_score: float + cvss_version: Literal["2.0", "3.0", "3.1"] + severity: Literal[1,2,3,4] = Field(default_factory=lambda: 1) + + model_config = { + "collection": "cves" + } + + @field_validator("cvss_score") + def validate_cvss_score(cls, v: float) -> float: + if v < 0.0 or v > 10.0: + raise ValueError("CVSS score must be between 0.0 and 10.0 inclusive") + return v + + def calculate_severity(self): + if self.cvss_version == "2.0": + if self.cvss_score == 10: + self.severity = 4 + elif self.cvss_score >= 7.0: + self.severity = 3 + elif self.cvss_score >= 4.0: + self.severity = 2 + else: + self.severity = 1 + else: # 3.0 or 3.1 + if self.cvss_score >= 9.0: + self.severity = 4 + elif self.cvss_score >= 7.0: + self.severity = 3 + elif self.cvss_score >= 4.0: + self.severity = 2 + else: + self.severity = 1 + + async def save(self, engine): + self.calculate_severity() + await engine.save(self) diff --git a/src/cyhy_db/utils/__init__.py b/src/cyhy_db/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/example/__init__.py b/src/example/__init__.py deleted file mode 100644 index 556a7d2..0000000 --- a/src/example/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -"""The example library.""" - -# We disable a Flake8 check for "Module imported but unused (F401)" here because -# although this import is not directly used, it populates the value -# package_name.__version__, which is used to get version information about this -# Python package. -from ._version import __version__ # noqa: F401 -from .example import example_div - -__all__ = ["example_div"] diff --git a/src/example/__main__.py b/src/example/__main__.py deleted file mode 100644 index 11a3238..0000000 --- a/src/example/__main__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Code to run if this package is used as a Python module.""" - -from .example import main - -main() diff --git a/src/example/_version.py b/src/example/_version.py deleted file mode 100644 index 3252c71..0000000 --- a/src/example/_version.py +++ /dev/null @@ -1,3 +0,0 @@ -"""This file defines the version of this module.""" - -__version__ = "0.2.1" diff --git a/src/example/data/secret.txt b/src/example/data/secret.txt deleted file mode 100644 index c40a49b..0000000 --- a/src/example/data/secret.txt +++ /dev/null @@ -1 +0,0 @@ -Three may keep a secret, if two of them are dead. diff --git a/src/example/example.py b/src/example/example.py deleted file mode 100644 index d3eda19..0000000 --- a/src/example/example.py +++ /dev/null @@ -1,103 +0,0 @@ -"""example is an example Python library and tool. - -Divide one integer by another and log the result. Also log some information -from an environment variable and a package resource. - -EXIT STATUS - This utility exits with one of the following values: - 0 Calculation completed successfully. - >0 An error occurred. - -Usage: - example [--log-level=LEVEL] - example (-h | --help) - -Options: - -h --help Show this message. - --log-level=LEVEL If specified, then the log level will be set to - the specified value. Valid values are "debug", "info", - "warning", "error", and "critical". [default: info] -""" - -# Standard Python Libraries -import logging -import os -import sys -from typing import Any, Dict - -# Third-Party Libraries -import docopt -import pkg_resources -from schema import And, Schema, SchemaError, Use - -from ._version import __version__ - -DEFAULT_ECHO_MESSAGE: str = "Hello World from the example default!" - - -def example_div(dividend: int, divisor: int) -> float: - """Print some logging messages.""" - logging.debug("This is a debug message") - logging.info("This is an info message") - logging.warning("This is a warning message") - logging.error("This is an error message") - logging.critical("This is a critical message") - return dividend / divisor - - -def main() -> None: - """Set up logging and call the example function.""" - args: Dict[str, str] = docopt.docopt(__doc__, version=__version__) - # Validate and convert arguments as needed - schema: Schema = Schema( - { - "--log-level": And( - str, - Use(str.lower), - lambda n: n in ("debug", "info", "warning", "error", "critical"), - error="Possible values for --log-level are " - + "debug, info, warning, error, and critical.", - ), - "": Use(int, error=" must be an integer."), - "": And( - Use(int), - lambda n: n != 0, - error=" must be an integer that is not 0.", - ), - str: object, # Don't care about other keys, if any - } - ) - - try: - validated_args: Dict[str, Any] = schema.validate(args) - except SchemaError as err: - # Exit because one or more of the arguments were invalid - print(err, file=sys.stderr) - sys.exit(1) - - # Assign validated arguments to variables - dividend: int = validated_args[""] - divisor: int = validated_args[""] - log_level: str = validated_args["--log-level"] - - # Set up logging - logging.basicConfig( - format="%(asctime)-15s %(levelname)s %(message)s", level=log_level.upper() - ) - - logging.info("%d / %d == %f", dividend, divisor, example_div(dividend, divisor)) - - # Access some data from an environment variable - message: str = os.getenv("ECHO_MESSAGE", DEFAULT_ECHO_MESSAGE) - logging.info('ECHO_MESSAGE="%s"', message) - - # Access some data from our package data (see the setup.py) - secret_message: str = ( - pkg_resources.resource_string("example", "data/secret.txt") - .decode("utf-8") - .strip() - ) - logging.info('Secret="%s"', secret_message) - - # Stop logging and clean up - logging.shutdown() diff --git a/tests/conftest.py b/tests/conftest.py index ba89c85..1582666 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,14 +4,106 @@ """ # Third-Party Libraries +import docker +from motor.motor_asyncio import AsyncIOMotorClient +from odmantic import AIOEngine import pytest +MONGO_INITDB_ROOT_USERNAME = os.environ.get("MONGO_INITDB_ROOT_USERNAME", "mongoadmin") +MONGO_INITDB_ROOT_PASSWORD = os.environ.get("MONGO_INITDB_ROOT_PASSWORD", "secret") +DATABASE_NAME = os.environ.get("DATABASE_NAME", "test") + +docker_client = docker.from_env() + + +@pytest.fixture(autouse=True) +def group_github_log_lines(request): + """Group log lines when running in GitHub actions.""" + # Group output from each test with workflow log groups + # https://help.github.com/en/actions/reference/workflow-commands-for-github-actions#grouping-log-lines + + if os.environ.get("GITHUB_ACTIONS") != "true": + # Not running in GitHub actions + yield + return + # Group using the current test name + print() + print(f"::group::{request.node.name}") + yield + print() + print("::endgroup::") + + +@pytest.fixture(scope="session") +def mongodb_container(mongo_image_tag): + """Fixture for the MongoDB test container.""" + container = docker_client.containers.run( + mongo_image_tag, + detach=True, + environment={ + "MONGO_INITDB_ROOT_USERNAME": MONGO_INITDB_ROOT_USERNAME, + "MONGO_INITDB_ROOT_PASSWORD": MONGO_INITDB_ROOT_PASSWORD, + }, + name="mongodb", + ports={"27017/tcp": None}, + volumes={}, + healthcheck={ + "test": ["CMD", "mongosh", "--eval", "'db.runCommand(\"ping\").ok'"], + "interval": 1000000000, # ns -> 1 second + "timeout": 1000000000, # ns -> 1 second + "retries": 5, + "start_period": 3000000000, # ns -> 3 seconds + }, + ) + TIMEOUT = 180 + # Wait for container to be healthy + for _ in range(TIMEOUT): + # Verify the container is still running + container.reload() + assert container.status == "running", "The container unexpectedly exited." + status = container.attrs["State"]["Health"]["Status"] + if status == "healthy": + break + time.sleep(1) + else: + assert ( + False + ), f"Container status did not transition to 'healthy' within {TIMEOUT} seconds." + + yield container + container.stop() + container.remove(force=True) + + +@pytest.fixture(scope="session") +def mongodb_engine(mongodb_container): + mongo_port = mongodb_container.attrs["NetworkSettings"]["Ports"]["27017/tcp"][0][ + "HostPort" + ] + mongo_uri = f"mongodb://{MONGO_INITDB_ROOT_USERNAME}:{MONGO_INITDB_ROOT_PASSWORD}@localhost:{mongo_port}" + + client = AsyncIOMotorClient(mongo_uri) + engine = AIOEngine(client=client, database=DATABASE_NAME) + return engine + def pytest_addoption(parser): """Add new commandline options to pytest.""" parser.addoption( "--runslow", action="store_true", default=False, help="run slow tests" ) + parser.addoption( + "--mongo-image-tag", + action="store", + default="docker.io/mongo:latest", + help="mongodb image tag to use for testing", + ) + + +@pytest.fixture(scope="session") +def mongo_image_tag(request): + """Get the image tag to test.""" + return request.config.getoption("--mongo-image-tag") def pytest_configure(config): diff --git a/tests/test_cve.py b/tests/test_cve.py new file mode 100644 index 0000000..6ae8a1b --- /dev/null +++ b/tests/test_cve.py @@ -0,0 +1,35 @@ +import pytest +from cyhy_db.models.cve import CVE + +severity_params = [ + ("2.0", 10, 4), + ("2.0", 7.0, 3), + ("2.0", 4.0, 2), + ("2.0", 0.0, 1), + ("3.0", 9.0, 4), + ("3.0", 7.0, 3), + ("3.0", 4.0, 2), + ("3.0", 0.0, 1), + ("3.1", 9.0, 4), + ("3.1", 7.0, 3), + ("3.1", 4.0, 2), + ("3.1", 0.0, 1), +] + +@pytest.mark.parametrize("version, score, expected_severity", severity_params) +def test_calculate_severity(version, score, expected_severity): + cve = CVE(cvss_version=version, cvss_score=score, id="test-cve") + cve.calculate_severity() + assert cve.severity == expected_severity, f"Failed for CVSS {version} with score {score}" + +def test_invalid_cvss_score(): + with pytest.raises(ValueError, match="CVSS score must be between 0.0 and 10.0 inclusive"): + CVE(cvss_version="3.1", cvss_score=11.0, id="test-cve") + +@pytest.mark.asyncio +async def test_save(mongodb_engine): + cve = CVE(cvss_version="3.1", cvss_score=9.0, id="test-cve") + await cve.save(mongodb_engine) + saved_cve = await mongodb_engine.find_one(CVE, CVE.id == "test-cve") + assert saved_cve is not None, "CVE not saved correctly" + assert saved_cve.severity == 4, "Severity not calculated correctly on save" diff --git a/tests/test_example.py b/tests/test_example.py deleted file mode 100644 index f8dea67..0000000 --- a/tests/test_example.py +++ /dev/null @@ -1,144 +0,0 @@ -#!/usr/bin/env pytest -vs -"""Tests for example.""" - -# Standard Python Libraries -import logging -import os -import sys -from unittest.mock import patch - -# Third-Party Libraries -import pytest - -# cisagov Libraries -import example - -div_params = [ - (1, 1, 1), - (2, 2, 1), - (0, 1, 0), - (8, 2, 4), -] - -log_levels = ( - "debug", - "info", - "warning", - "error", - "critical", -) - -# define sources of version strings -RELEASE_TAG = os.getenv("RELEASE_TAG") -PROJECT_VERSION = example.__version__ - - -def test_stdout_version(capsys): - """Verify that version string sent to stdout agrees with the module version.""" - with pytest.raises(SystemExit): - with patch.object(sys, "argv", ["bogus", "--version"]): - example.example.main() - captured = capsys.readouterr() - assert ( - captured.out == f"{PROJECT_VERSION}\n" - ), "standard output by '--version' should agree with module.__version__" - - -def test_running_as_module(capsys): - """Verify that the __main__.py file loads correctly.""" - with pytest.raises(SystemExit): - with patch.object(sys, "argv", ["bogus", "--version"]): - # F401 is a "Module imported but unused" warning. This import - # emulates how this project would be run as a module. The only thing - # being done by __main__ is importing the main entrypoint of the - # package and running it, so there is nothing to use from this - # import. As a result, we can safely ignore this warning. - # cisagov Libraries - import example.__main__ # noqa: F401 - captured = capsys.readouterr() - assert ( - captured.out == f"{PROJECT_VERSION}\n" - ), "standard output by '--version' should agree with module.__version__" - - -@pytest.mark.skipif( - RELEASE_TAG in [None, ""], reason="this is not a release (RELEASE_TAG not set)" -) -def test_release_version(): - """Verify that release tag version agrees with the module version.""" - assert ( - RELEASE_TAG == f"v{PROJECT_VERSION}" - ), "RELEASE_TAG does not match the project version" - - -@pytest.mark.parametrize("level", log_levels) -def test_log_levels(level): - """Validate commandline log-level arguments.""" - with patch.object(sys, "argv", ["bogus", f"--log-level={level}", "1", "1"]): - with patch.object(logging.root, "handlers", []): - assert ( - logging.root.hasHandlers() is False - ), "root logger should not have handlers yet" - return_code = None - try: - example.example.main() - except SystemExit as sys_exit: - return_code = sys_exit.code - assert return_code is None, "main() should return success" - assert ( - logging.root.hasHandlers() is True - ), "root logger should now have a handler" - assert ( - logging.getLevelName(logging.root.getEffectiveLevel()) == level.upper() - ), f"root logger level should be set to {level.upper()}" - assert return_code is None, "main() should return success" - - -def test_bad_log_level(): - """Validate bad log-level argument returns error.""" - with patch.object(sys, "argv", ["bogus", "--log-level=emergency", "1", "1"]): - return_code = None - try: - example.example.main() - except SystemExit as sys_exit: - return_code = sys_exit.code - assert return_code == 1, "main() should exit with error" - - -@pytest.mark.parametrize("dividend, divisor, quotient", div_params) -def test_division(dividend, divisor, quotient): - """Verify division results.""" - result = example.example_div(dividend, divisor) - assert result == quotient, "result should equal quotient" - - -@pytest.mark.slow -def test_slow_division(): - """Example of using a custom marker. - - This test will only be run if --runslow is passed to pytest. - Look in conftest.py to see how this is implemented. - """ - # Standard Python Libraries - import time - - result = example.example_div(256, 16) - time.sleep(4) - assert result == 16, "result should equal be 16" - - -def test_zero_division(): - """Verify that division by zero throws the correct exception.""" - with pytest.raises(ZeroDivisionError): - example.example_div(1, 0) - - -def test_zero_divisor_argument(): - """Verify that a divisor of zero is handled as expected.""" - with patch.object(sys, "argv", ["bogus", "1", "0"]): - return_code = None - try: - example.example.main() - except SystemExit as sys_exit: - return_code = sys_exit.code - assert return_code == 1, "main() should exit with error" From 904692bd586146abf50b5578f6ac8c4901863a14 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 24 Jan 2024 16:27:27 -0500 Subject: [PATCH 002/139] Apply changes made by linters Add comments to severity calculation. --- src/cyhy_db/_version.py | 2 +- src/cyhy_db/models/cve.py | 33 +++++++++++++++++++++++++-------- tests/test_cve.py | 24 +++++++++++++++++------- 3 files changed, 43 insertions(+), 16 deletions(-) diff --git a/src/cyhy_db/_version.py b/src/cyhy_db/_version.py index b3c06d4..f102a9c 100644 --- a/src/cyhy_db/_version.py +++ b/src/cyhy_db/_version.py @@ -1 +1 @@ -__version__ = "0.0.1" \ No newline at end of file +__version__ = "0.0.1" diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index f8dfcb8..c956e3f 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -1,24 +1,41 @@ -from odmantic import Model, Field -from pydantic import ValidationInfo, field_validator +# Standard Python Libraries from typing import Literal +# Third-Party Libraries +from odmantic import Field, Model +from pydantic import ValidationInfo, field_validator + + class CVE(Model): id: str = Field(primary_field=True) cvss_score: float cvss_version: Literal["2.0", "3.0", "3.1"] - severity: Literal[1,2,3,4] = Field(default_factory=lambda: 1) + severity: Literal[1, 2, 3, 4] = Field(default_factory=lambda: 1) - model_config = { - "collection": "cves" - } + model_config = {"collection": "cves"} @field_validator("cvss_score") def validate_cvss_score(cls, v: float) -> float: if v < 0.0 or v > 10.0: raise ValueError("CVSS score must be between 0.0 and 10.0 inclusive") return v - + def calculate_severity(self): + # Calculate severity from cvss on save + # Source: https://nvd.nist.gov/vuln-metrics/cvss + # + # Notes: + # - The CVSS score to severity mapping is not continuous (e.g. a + # score of 8.95 is undefined according to their table). However, + # the CVSS equation documentation + # (https://www.first.org/cvss/specification-document#CVSS-v3-1-Equations) + # specifies that all CVSS scores are rounded up to the nearest tenth + # of a point, so our severity mapping below is valid. + # - CVSSv3 specifies that a score of 0.0 has a severity of "None", but + # we have chosen to map 0.0 to severity 1 ("Low") because CyHy code + # has historically assumed severities between 1 and 4 (inclusive). + # Since we have not seen CVSSv3 scores lower than 3.1, this will + # hopefully never be an issue. if self.cvss_version == "2.0": if self.cvss_score == 10: self.severity = 4 @@ -28,7 +45,7 @@ def calculate_severity(self): self.severity = 2 else: self.severity = 1 - else: # 3.0 or 3.1 + else: # 3.0 or 3.1 if self.cvss_score >= 9.0: self.severity = 4 elif self.cvss_score >= 7.0: diff --git a/tests/test_cve.py b/tests/test_cve.py index 6ae8a1b..660298d 100644 --- a/tests/test_cve.py +++ b/tests/test_cve.py @@ -1,4 +1,7 @@ +# Third-Party Libraries import pytest + +# cisagov Libraries from cyhy_db.models.cve import CVE severity_params = [ @@ -16,20 +19,27 @@ ("3.1", 0.0, 1), ] + @pytest.mark.parametrize("version, score, expected_severity", severity_params) def test_calculate_severity(version, score, expected_severity): cve = CVE(cvss_version=version, cvss_score=score, id="test-cve") cve.calculate_severity() - assert cve.severity == expected_severity, f"Failed for CVSS {version} with score {score}" + assert ( + cve.severity == expected_severity + ), f"Failed for CVSS {version} with score {score}" + def test_invalid_cvss_score(): - with pytest.raises(ValueError, match="CVSS score must be between 0.0 and 10.0 inclusive"): + with pytest.raises( + ValueError, match="CVSS score must be between 0.0 and 10.0 inclusive" + ): CVE(cvss_version="3.1", cvss_score=11.0, id="test-cve") + @pytest.mark.asyncio async def test_save(mongodb_engine): - cve = CVE(cvss_version="3.1", cvss_score=9.0, id="test-cve") - await cve.save(mongodb_engine) - saved_cve = await mongodb_engine.find_one(CVE, CVE.id == "test-cve") - assert saved_cve is not None, "CVE not saved correctly" - assert saved_cve.severity == 4, "Severity not calculated correctly on save" + cve = CVE(cvss_version="3.1", cvss_score=9.0, id="test-cve") + await cve.save(mongodb_engine) + saved_cve = await mongodb_engine.find_one(CVE, CVE.id == "test-cve") + assert saved_cve is not None, "CVE not saved correctly" + assert saved_cve.severity == 4, "Severity not calculated correctly on save" From 6593b87f27226ef236d38e1a8b932f7a0dc1bb65 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 24 Jan 2024 16:35:30 -0500 Subject: [PATCH 003/139] Add doc strings where needed --- src/cyhy_db/models/cve.py | 9 ++++++++- src/cyhy_db/utils/__init__.py | 1 + tests/conftest.py | 1 + tests/test_cve.py | 5 +++++ 4 files changed, 15 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index c956e3f..417c4f1 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -1,12 +1,16 @@ +"""This module defines the CVE model.""" + # Standard Python Libraries from typing import Literal # Third-Party Libraries from odmantic import Field, Model -from pydantic import ValidationInfo, field_validator +from pydantic import field_validator class CVE(Model): + """This class represents the CVE model.""" + id: str = Field(primary_field=True) cvss_score: float cvss_version: Literal["2.0", "3.0", "3.1"] @@ -16,11 +20,13 @@ class CVE(Model): @field_validator("cvss_score") def validate_cvss_score(cls, v: float) -> float: + """Validate the CVSS score.""" if v < 0.0 or v > 10.0: raise ValueError("CVSS score must be between 0.0 and 10.0 inclusive") return v def calculate_severity(self): + """Calculate the severity from the CVSS score.""" # Calculate severity from cvss on save # Source: https://nvd.nist.gov/vuln-metrics/cvss # @@ -56,5 +62,6 @@ def calculate_severity(self): self.severity = 1 async def save(self, engine): + """Save the CVE to the database.""" self.calculate_severity() await engine.save(self) diff --git a/src/cyhy_db/utils/__init__.py b/src/cyhy_db/utils/__init__.py index e69de29..c12b5c9 100644 --- a/src/cyhy_db/utils/__init__.py +++ b/src/cyhy_db/utils/__init__.py @@ -0,0 +1 @@ +"""Utility functions for cyhy_db.""" diff --git a/tests/conftest.py b/tests/conftest.py index 1582666..51319b0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -77,6 +77,7 @@ def mongodb_container(mongo_image_tag): @pytest.fixture(scope="session") def mongodb_engine(mongodb_container): + """Fixture for the MongoDB engine.""" mongo_port = mongodb_container.attrs["NetworkSettings"]["Ports"]["27017/tcp"][0][ "HostPort" ] diff --git a/tests/test_cve.py b/tests/test_cve.py index 660298d..27380c6 100644 --- a/tests/test_cve.py +++ b/tests/test_cve.py @@ -1,3 +1,5 @@ +"""Test CVE model functionality.""" + # Third-Party Libraries import pytest @@ -22,6 +24,7 @@ @pytest.mark.parametrize("version, score, expected_severity", severity_params) def test_calculate_severity(version, score, expected_severity): + """Test that the severity is calculated correctly.""" cve = CVE(cvss_version=version, cvss_score=score, id="test-cve") cve.calculate_severity() assert ( @@ -30,6 +33,7 @@ def test_calculate_severity(version, score, expected_severity): def test_invalid_cvss_score(): + """Test that an invalid CVSS score raises a ValueError.""" with pytest.raises( ValueError, match="CVSS score must be between 0.0 and 10.0 inclusive" ): @@ -38,6 +42,7 @@ def test_invalid_cvss_score(): @pytest.mark.asyncio async def test_save(mongodb_engine): + """Test that the severity is calculated correctly on save.""" cve = CVE(cvss_version="3.1", cvss_score=9.0, id="test-cve") await cve.save(mongodb_engine) saved_cve = await mongodb_engine.find_one(CVE, CVE.id == "test-cve") From 312752be62406265ab4878c6e3b7ded265fbb205 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 24 Jan 2024 17:44:35 -0500 Subject: [PATCH 004/139] Simplify field validation --- src/cyhy_db/models/cve.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 417c4f1..1d365e5 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -5,26 +5,18 @@ # Third-Party Libraries from odmantic import Field, Model -from pydantic import field_validator class CVE(Model): """This class represents the CVE model.""" id: str = Field(primary_field=True) - cvss_score: float + cvss_score: float = Field(ge=0.0, le=10.0) cvss_version: Literal["2.0", "3.0", "3.1"] - severity: Literal[1, 2, 3, 4] = Field(default_factory=lambda: 1) + severity: Literal[1, 2, 3, 4] = 1 model_config = {"collection": "cves"} - @field_validator("cvss_score") - def validate_cvss_score(cls, v: float) -> float: - """Validate the CVSS score.""" - if v < 0.0 or v > 10.0: - raise ValueError("CVSS score must be between 0.0 and 10.0 inclusive") - return v - def calculate_severity(self): """Calculate the severity from the CVSS score.""" # Calculate severity from cvss on save From 250d6f669cef490ccee76a6cc80f0027633e95aa Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 24 Jan 2024 17:44:53 -0500 Subject: [PATCH 005/139] Add test case for invalid CVSS score --- tests/test_cve.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_cve.py b/tests/test_cve.py index 27380c6..eea5f0f 100644 --- a/tests/test_cve.py +++ b/tests/test_cve.py @@ -22,6 +22,7 @@ ] + @pytest.mark.parametrize("version, score, expected_severity", severity_params) def test_calculate_severity(version, score, expected_severity): """Test that the severity is calculated correctly.""" @@ -31,14 +32,13 @@ def test_calculate_severity(version, score, expected_severity): cve.severity == expected_severity ), f"Failed for CVSS {version} with score {score}" - -def test_invalid_cvss_score(): +@pytest.mark.parametrize("bad_score", [-1.0, 11.0]) +def test_invalid_cvss_score(bad_score): """Test that an invalid CVSS score raises a ValueError.""" with pytest.raises( - ValueError, match="CVSS score must be between 0.0 and 10.0 inclusive" + ValueError, match=f"than or equal to" ): - CVE(cvss_version="3.1", cvss_score=11.0, id="test-cve") - + CVE(cvss_version="3.1", cvss_score=bad_score, id="test-cve") @pytest.mark.asyncio async def test_save(mongodb_engine): From 3f39aae747d82bd9a5be13122880ee65b51340cc Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 26 Jan 2024 17:39:46 -0500 Subject: [PATCH 006/139] Switch to MongoEngine --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 00149ef..e2de6f5 100644 --- a/setup.py +++ b/setup.py @@ -92,7 +92,7 @@ def get_version(version_file): include_package_data=True, install_requires=[ "docopt == 0.6.2", - "odmantic == 1.0.0", + "mongoengine == 0.27.0", "schema == 0.7.5", "setuptools >= 69.0.3", ], From 8999e2666129431efea414735eeb3efaf9c78f0c Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 26 Jan 2024 17:40:55 -0500 Subject: [PATCH 007/139] Add utcnow function to utils module --- src/cyhy_db/utils/__init__.py | 4 ++++ src/cyhy_db/utils/time.py | 6 ++++++ 2 files changed, 10 insertions(+) create mode 100644 src/cyhy_db/utils/time.py diff --git a/src/cyhy_db/utils/__init__.py b/src/cyhy_db/utils/__init__.py index c12b5c9..a9bfaf4 100644 --- a/src/cyhy_db/utils/__init__.py +++ b/src/cyhy_db/utils/__init__.py @@ -1 +1,5 @@ """Utility functions for cyhy_db.""" + +from .time import utcnow + +__all__ = ["utcnow"] diff --git a/src/cyhy_db/utils/time.py b/src/cyhy_db/utils/time.py new file mode 100644 index 0000000..08ad3ce --- /dev/null +++ b/src/cyhy_db/utils/time.py @@ -0,0 +1,6 @@ +from datetime import datetime, timezone + + +def utcnow() -> datetime: + """Returns a timezone-aware datetime object with the current time in UTC.""" + return datetime.now(timezone.utc) From 3a291cd7ecd59c8c0d92e7a69286166d3ddda424 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 26 Jan 2024 17:41:16 -0500 Subject: [PATCH 008/139] Update MongoDB connection and add support for Mongo Express --- tests/conftest.py | 58 +++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 51 insertions(+), 7 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 51319b0..85b69c2 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -5,8 +5,7 @@ # Third-Party Libraries import docker -from motor.motor_asyncio import AsyncIOMotorClient -from odmantic import AIOEngine +from mongoengine import connect, disconnect import pytest MONGO_INITDB_ROOT_USERNAME = os.environ.get("MONGO_INITDB_ROOT_USERNAME", "mongoadmin") @@ -76,16 +75,55 @@ def mongodb_container(mongo_image_tag): @pytest.fixture(scope="session") -def mongodb_engine(mongodb_container): - """Fixture for the MongoDB engine.""" +def mongo_express_container(mongodb_container, request): + if not request.config.getoption("--mongo-express"): + yield None + return + + # Configuration for Mongo Express + mongo_express_container = docker_client.containers.run( + "mongo-express", + environment={ + "ME_CONFIG_MONGODB_ADMINUSERNAME": MONGO_INITDB_ROOT_USERNAME, + "ME_CONFIG_MONGODB_ADMINPASSWORD": MONGO_INITDB_ROOT_PASSWORD, + "ME_CONFIG_MONGODB_SERVER": "mongodb", + "ME_CONFIG_MONGODB_ENABLE_ADMIN": "true", + }, + links={"mongodb": "mongodb"}, + ports={"8081/tcp": 8081}, + detach=True, + ) + + def fin(): + if request.config.getoption("--mongo-express"): + input("\n\nPress Enter to stop Mongo Express and MongoDB containers...") + mongo_express_container.stop() + mongo_express_container.remove(force=True) + + request.addfinalizer(fin) + yield mongo_express_container + + +@pytest.fixture(scope="session") +def mongo_uri(mongodb_container): + """Fixture for the MongoDB URI.""" mongo_port = mongodb_container.attrs["NetworkSettings"]["Ports"]["27017/tcp"][0][ "HostPort" ] mongo_uri = f"mongodb://{MONGO_INITDB_ROOT_USERNAME}:{MONGO_INITDB_ROOT_PASSWORD}@localhost:{mongo_port}" + yield mongo_uri + + +@pytest.fixture(scope="session") +def mongodb_engine(mongo_uri, mongo_express_container): + """Fixture for the MongoDB engine.""" + connect( + host=mongo_uri, alias="default", db=DATABASE_NAME, uuidRepresentation="standard" + ) + + yield - client = AsyncIOMotorClient(mongo_uri) - engine = AIOEngine(client=client, database=DATABASE_NAME) - return engine + disconnect() def pytest_addoption(parser): @@ -99,6 +137,12 @@ def pytest_addoption(parser): default="docker.io/mongo:latest", help="mongodb image tag to use for testing", ) + parser.addoption( + "--mongo-express", + action="store_true", + default=False, + help="run Mongo Express for database inspection", + ) @pytest.fixture(scope="session") From 97359b94bc7688e11b9e75ffc4a5c76faa5e1235 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 26 Jan 2024 17:41:55 -0500 Subject: [PATCH 009/139] Refactor CVE model and test This commit refactors the CVE model and test in the `cyhy_db` package. The changes include: - Replacing the `odmantic` library with `mongoengine` for the model definition. - Updating the field types and constraints in the CVE model. - Removing the `calculate_severity` method and integrating the severity calculation directly in the `save` method. - Updating the test cases to reflect the changes in the model. --- src/cyhy_db/models/cve.py | 44 +++++++++------------------------------ tests/test_cve.py | 19 +++++++++-------- 2 files changed, 20 insertions(+), 43 deletions(-) diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 1d365e5..2599677 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -1,39 +1,16 @@ -"""This module defines the CVE model.""" - -# Standard Python Libraries -from typing import Literal - # Third-Party Libraries -from odmantic import Field, Model - +from mongoengine import Document, FloatField, IntField, StringField -class CVE(Model): - """This class represents the CVE model.""" - id: str = Field(primary_field=True) - cvss_score: float = Field(ge=0.0, le=10.0) - cvss_version: Literal["2.0", "3.0", "3.1"] - severity: Literal[1, 2, 3, 4] = 1 +class CVE(Document): + id = StringField(primary_key=True) + cvss_score = FloatField(min_value=0.0, max_value=10.0) + cvss_version = StringField(choices=["2.0", "3.0", "3.1"]) + severity = IntField(choices=[1, 2, 3, 4], default=1) - model_config = {"collection": "cves"} + meta = {"collection": "cves"} def calculate_severity(self): - """Calculate the severity from the CVSS score.""" - # Calculate severity from cvss on save - # Source: https://nvd.nist.gov/vuln-metrics/cvss - # - # Notes: - # - The CVSS score to severity mapping is not continuous (e.g. a - # score of 8.95 is undefined according to their table). However, - # the CVSS equation documentation - # (https://www.first.org/cvss/specification-document#CVSS-v3-1-Equations) - # specifies that all CVSS scores are rounded up to the nearest tenth - # of a point, so our severity mapping below is valid. - # - CVSSv3 specifies that a score of 0.0 has a severity of "None", but - # we have chosen to map 0.0 to severity 1 ("Low") because CyHy code - # has historically assumed severities between 1 and 4 (inclusive). - # Since we have not seen CVSSv3 scores lower than 3.1, this will - # hopefully never be an issue. if self.cvss_version == "2.0": if self.cvss_score == 10: self.severity = 4 @@ -43,7 +20,7 @@ def calculate_severity(self): self.severity = 2 else: self.severity = 1 - else: # 3.0 or 3.1 + else: # CVSS versions 3.0 or 3.1 if self.cvss_score >= 9.0: self.severity = 4 elif self.cvss_score >= 7.0: @@ -53,7 +30,6 @@ def calculate_severity(self): else: self.severity = 1 - async def save(self, engine): - """Save the CVE to the database.""" + def save(self, *args, **kwargs): self.calculate_severity() - await engine.save(self) + return super().save(*args, **kwargs) diff --git a/tests/test_cve.py b/tests/test_cve.py index eea5f0f..51ba990 100644 --- a/tests/test_cve.py +++ b/tests/test_cve.py @@ -1,6 +1,7 @@ """Test CVE model functionality.""" # Third-Party Libraries +from mongoengine import ValidationError import pytest # cisagov Libraries @@ -22,7 +23,6 @@ ] - @pytest.mark.parametrize("version, score, expected_severity", severity_params) def test_calculate_severity(version, score, expected_severity): """Test that the severity is calculated correctly.""" @@ -32,19 +32,20 @@ def test_calculate_severity(version, score, expected_severity): cve.severity == expected_severity ), f"Failed for CVSS {version} with score {score}" + @pytest.mark.parametrize("bad_score", [-1.0, 11.0]) def test_invalid_cvss_score(bad_score): """Test that an invalid CVSS score raises a ValueError.""" - with pytest.raises( - ValueError, match=f"than or equal to" - ): - CVE(cvss_version="3.1", cvss_score=bad_score, id="test-cve") + cve = CVE(cvss_version="3.1", cvss_score=bad_score, id="test-cve") + with pytest.raises(ValidationError): + cve.validate() # Explicitly call validate to trigger validation -@pytest.mark.asyncio -async def test_save(mongodb_engine): + +def test_save(mongodb_engine): """Test that the severity is calculated correctly on save.""" cve = CVE(cvss_version="3.1", cvss_score=9.0, id="test-cve") - await cve.save(mongodb_engine) - saved_cve = await mongodb_engine.find_one(CVE, CVE.id == "test-cve") + cve.save() # Saving the object + saved_cve = CVE.objects(id="test-cve").first() # Retrieving the object + assert saved_cve is not None, "CVE not saved correctly" assert saved_cve.severity == 4, "Severity not calculated correctly on save" From a5684549a6d094d2d599205e765e541a5726b4ac Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 26 Jan 2024 17:42:12 -0500 Subject: [PATCH 010/139] Add IPAddressField and test cases for IP address validation and conversion --- src/cyhy_db/models/ip_address.py | 30 +++++++++++++++++ tests/test_ip_address.py | 55 ++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) create mode 100644 src/cyhy_db/models/ip_address.py create mode 100644 tests/test_ip_address.py diff --git a/src/cyhy_db/models/ip_address.py b/src/cyhy_db/models/ip_address.py new file mode 100644 index 0000000..bfaa277 --- /dev/null +++ b/src/cyhy_db/models/ip_address.py @@ -0,0 +1,30 @@ +import ipaddress +from mongoengine.base.fields import BaseField + + +class IPAddressField(BaseField): + def validate(self, value): + if not isinstance(value, (ipaddress.IPv4Address, ipaddress.IPv6Address)): + self.error(f'Value "{value}" is not a valid IP address.') + + def to_mongo(self, value): + if value is None: + return value + if isinstance(value, (ipaddress.IPv4Address, ipaddress.IPv6Address)): + return str(value) + self.error(f'Value "{value}" cannot be converted to a string IP address.') + + def prepare_query_value(self, op, value): + return super().prepare_query_value(op, self.to_mongo(value)) + + @staticmethod + def _parse_ip(value): + if isinstance(value, str): + return ipaddress.ip_address(value.strip()) + return value + + def to_python(self, value): + try: + return self._parse_ip(value) + except ValueError as e: + self.error(f'Value "{value}" is not a valid IP address.') diff --git a/tests/test_ip_address.py b/tests/test_ip_address.py new file mode 100644 index 0000000..efc3b3d --- /dev/null +++ b/tests/test_ip_address.py @@ -0,0 +1,55 @@ +import pytest +from cyhy_db.models.ip_address import IPAddressField +from mongoengine import Document, ValidationError +import ipaddress + + +class HasIpDocument(Document): + ip = IPAddressField() + + +def test_ip_address_type(): + doc = HasIpDocument(ip="1.2.3.4") + assert isinstance(doc.ip, (ipaddress.IPv4Address, ipaddress.IPv6Address)) + + +@pytest.mark.parametrize( + "valid_ip", + [ + "192.168.1.1", + "10.0.0.1", + "255.255.255.255", + ], +) +def test_valid_ip_address_field(valid_ip): + try: + HasIpDocument(ip=valid_ip).validate() + except ValidationError: + pytest.fail(f"Valid IP address {valid_ip} raised ValidationError") + + +@pytest.mark.parametrize( + "invalid_ip", + [ + "256.256.256.256", + "123.456.789.0", + "abc.def.ghi.jkl", + ], +) +def test_invalid_ip_address_field(invalid_ip): + with pytest.raises(ValidationError): + HasIpDocument(ip=invalid_ip).validate() + + +def test_save_ip_address(mongodb_engine): + test_document = HasIpDocument(ip=ipaddress.IPv4Address("192.168.1.1")) + test_document.save() + assert test_document.id is not None, "TestDocument instance was not saved correctly" + + +def test_retrieve_ip_address(mongodb_engine): + retrieved_doc = HasIpDocument.objects().first() + assert retrieved_doc is not None, "TestDocument instance was not retrieved" + assert retrieved_doc.ip == ipaddress.IPv4Address( + "192.168.1.1" + ), "Retrieved IP address does not match the saved IP address" From 9c97a24694b3f7a5d6847ddcd310367c9d22f942 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 26 Jan 2024 17:42:32 -0500 Subject: [PATCH 011/139] Add ScanDoc and SnapshotDoc models --- src/cyhy_db/models/scan_doc.py | 70 ++++++++++++++++++++++++++++++ src/cyhy_db/models/snapshot_doc.py | 7 +++ 2 files changed, 77 insertions(+) create mode 100644 src/cyhy_db/models/scan_doc.py create mode 100644 src/cyhy_db/models/snapshot_doc.py diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py new file mode 100644 index 0000000..3ac705e --- /dev/null +++ b/src/cyhy_db/models/scan_doc.py @@ -0,0 +1,70 @@ +from mongoengine import ( + Document, + StringField, + IntField, + BooleanField, + ListField, + ReferenceField, + DateTimeField, +) +import datetime +from .ip_address import IPAddressField +from collections.abc import Iterable + + +class ScanDoc(Document): + _ip = IPAddressField(db_field="ip", required=True) + ip_int = IntField(required=True) + latest = BooleanField(default=True) + owner = StringField(required=True) + snapshots = ListField(ReferenceField("SnapshotDoc")) + source = StringField(required=True) + time = DateTimeField(default=datetime.datetime.utcnow, required=True) + + meta = { + "indexes": [ + {"fields": ["latest", "ip_int"], "unique": False}, + {"fields": ["time", "owner"], "unique": False}, + {"fields": ["ip_int"], "unique": False}, + {"fields": ["snapshots"], "unique": False, "sparse": True}, + ] + } + + @property + def ip(self): + return self._ip + + @ip.setter + def ip(self, new_ip): + self._ip = new_ip + # Convert IP to an integer and store in ip_int + try: + self.ip_int = int(self._ip) + except ValueError: + raise ValueError(f"Invalid IP address: {new_ip}") + + # Custom methods + def reset_latest_flag_by_owner(self, owner): + ScanDoc.objects(latest=True, owner=owner).update(latest=False) + + def reset_latest_flag_by_ip(self, ips): + ip_ints = [int(x) for x in ips] if isinstance(ips, Iterable) else [int(ips)] + ScanDoc.objects(latest=True, ip_int__in=ip_ints).update(latest=False) + + def tag_latest(self, owners, snapshot_oid): + ScanDoc.objects(latest=True, owner__in=owners).update( + push__snapshots=snapshot_oid + ) + + def tag_matching(self, existing_snapshot_oids, new_snapshot_oid): + ScanDoc.objects(snapshots__in=existing_snapshot_oids).update( + push__snapshots=new_snapshot_oid + ) + + def tag_timespan(self, owner, snapshot_oid, start_time, end_time): + ScanDoc.objects(time__gte=start_time, time__lte=end_time, owner=owner).update( + push__snapshots=snapshot_oid + ) + + def remove_tag(self, snapshot_oid): + ScanDoc.objects(snapshots=snapshot_oid).update(pull__snapshots=snapshot_oid) diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py new file mode 100644 index 0000000..b1e9ecc --- /dev/null +++ b/src/cyhy_db/models/snapshot_doc.py @@ -0,0 +1,7 @@ +from mongoengine import ( + Document, +) + + +class SnapshotDoc(Document): + pass From 09fd02c1994e0f5b35d4679d1bc5c4f11025f05f Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 26 Jan 2024 17:43:47 -0500 Subject: [PATCH 012/139] Sort imports --- src/cyhy_db/models/ip_address.py | 3 +++ src/cyhy_db/models/scan_doc.py | 14 +++++++++----- src/cyhy_db/models/snapshot_doc.py | 5 ++--- src/cyhy_db/utils/time.py | 1 + tests/test_ip_address.py | 9 +++++++-- 5 files changed, 22 insertions(+), 10 deletions(-) diff --git a/src/cyhy_db/models/ip_address.py b/src/cyhy_db/models/ip_address.py index bfaa277..4974fef 100644 --- a/src/cyhy_db/models/ip_address.py +++ b/src/cyhy_db/models/ip_address.py @@ -1,4 +1,7 @@ +# Standard Python Libraries import ipaddress + +# Third-Party Libraries from mongoengine.base.fields import BaseField diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 3ac705e..9c761d6 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -1,15 +1,19 @@ +# Standard Python Libraries +from collections.abc import Iterable +import datetime + +# Third-Party Libraries from mongoengine import ( + BooleanField, + DateTimeField, Document, - StringField, IntField, - BooleanField, ListField, ReferenceField, - DateTimeField, + StringField, ) -import datetime + from .ip_address import IPAddressField -from collections.abc import Iterable class ScanDoc(Document): diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index b1e9ecc..72ba359 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -1,6 +1,5 @@ -from mongoengine import ( - Document, -) +# Third-Party Libraries +from mongoengine import Document class SnapshotDoc(Document): diff --git a/src/cyhy_db/utils/time.py b/src/cyhy_db/utils/time.py index 08ad3ce..c3e4512 100644 --- a/src/cyhy_db/utils/time.py +++ b/src/cyhy_db/utils/time.py @@ -1,3 +1,4 @@ +# Standard Python Libraries from datetime import datetime, timezone diff --git a/tests/test_ip_address.py b/tests/test_ip_address.py index efc3b3d..893b381 100644 --- a/tests/test_ip_address.py +++ b/tests/test_ip_address.py @@ -1,7 +1,12 @@ +# Standard Python Libraries +import ipaddress + +# Third-Party Libraries +from mongoengine import Document, ValidationError import pytest + +# cisagov Libraries from cyhy_db.models.ip_address import IPAddressField -from mongoengine import Document, ValidationError -import ipaddress class HasIpDocument(Document): From c29996e8e20a9c87b4ac088e8cd78d8ad02b4e3d Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:45:35 -0500 Subject: [PATCH 013/139] Implement CVE, ScanDoc with beanie --- src/cyhy_db/models/__init__.py | 5 ++ src/cyhy_db/models/cve.py | 19 +++-- src/cyhy_db/models/ip_address.py | 33 --------- src/cyhy_db/models/scan_doc.py | 108 +++++++++++++---------------- src/cyhy_db/models/snapshot_doc.py | 2 +- 5 files changed, 65 insertions(+), 102 deletions(-) create mode 100644 src/cyhy_db/models/__init__.py delete mode 100644 src/cyhy_db/models/ip_address.py diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py new file mode 100644 index 0000000..87d45e9 --- /dev/null +++ b/src/cyhy_db/models/__init__.py @@ -0,0 +1,5 @@ +from .cve import CVE +from .scan_doc import ScanDoc +from .snapshot_doc import SnapshotDoc + +__all__ = ["CVE", "ScanDoc", "SnapshotDoc"] diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 2599677..54828c8 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -1,15 +1,18 @@ # Third-Party Libraries -from mongoengine import Document, FloatField, IntField, StringField +from beanie import Document, Indexed, ValidateOnSave, before_event +from pydantic import Field class CVE(Document): - id = StringField(primary_key=True) - cvss_score = FloatField(min_value=0.0, max_value=10.0) - cvss_version = StringField(choices=["2.0", "3.0", "3.1"]) - severity = IntField(choices=[1, 2, 3, 4], default=1) + id: str = Indexed(primary_field=True) # CVE ID + cvss_score: float = Field(ge=0.0, le=10.0) + cvss_version: str = Field(enum=["2.0", "3.0", "3.1"]) + severity: int = Field(ge=1, le=4, default=1) - meta = {"collection": "cves"} + class Settings: + name = "cves" + @before_event(ValidateOnSave) def calculate_severity(self): if self.cvss_version == "2.0": if self.cvss_score == 10: @@ -29,7 +32,3 @@ def calculate_severity(self): self.severity = 2 else: self.severity = 1 - - def save(self, *args, **kwargs): - self.calculate_severity() - return super().save(*args, **kwargs) diff --git a/src/cyhy_db/models/ip_address.py b/src/cyhy_db/models/ip_address.py deleted file mode 100644 index 4974fef..0000000 --- a/src/cyhy_db/models/ip_address.py +++ /dev/null @@ -1,33 +0,0 @@ -# Standard Python Libraries -import ipaddress - -# Third-Party Libraries -from mongoengine.base.fields import BaseField - - -class IPAddressField(BaseField): - def validate(self, value): - if not isinstance(value, (ipaddress.IPv4Address, ipaddress.IPv6Address)): - self.error(f'Value "{value}" is not a valid IP address.') - - def to_mongo(self, value): - if value is None: - return value - if isinstance(value, (ipaddress.IPv4Address, ipaddress.IPv6Address)): - return str(value) - self.error(f'Value "{value}" cannot be converted to a string IP address.') - - def prepare_query_value(self, op, value): - return super().prepare_query_value(op, self.to_mongo(value)) - - @staticmethod - def _parse_ip(value): - if isinstance(value, str): - return ipaddress.ip_address(value.strip()) - return value - - def to_python(self, value): - try: - return self._parse_ip(value) - except ValueError as e: - self.error(f'Value "{value}" is not a valid IP address.') diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 9c761d6..6d3af58 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -1,74 +1,66 @@ # Standard Python Libraries -from collections.abc import Iterable import datetime +import ipaddress +from typing import Any, Dict, Iterable # Third-Party Libraries -from mongoengine import ( - BooleanField, - DateTimeField, - Document, - IntField, - ListField, - ReferenceField, - StringField, -) - -from .ip_address import IPAddressField +from beanie import Document, Link +from beanie.operators import Push, Set +from pydantic import Field, model_validator +from pymongo import ASCENDING, IndexModel class ScanDoc(Document): - _ip = IPAddressField(db_field="ip", required=True) - ip_int = IntField(required=True) - latest = BooleanField(default=True) - owner = StringField(required=True) - snapshots = ListField(ReferenceField("SnapshotDoc")) - source = StringField(required=True) - time = DateTimeField(default=datetime.datetime.utcnow, required=True) - - meta = { - "indexes": [ - {"fields": ["latest", "ip_int"], "unique": False}, - {"fields": ["time", "owner"], "unique": False}, - {"fields": ["ip_int"], "unique": False}, - {"fields": ["snapshots"], "unique": False, "sparse": True}, - ] - } - - @property - def ip(self): - return self._ip + ip: ipaddress.IPv4Address = Field(...) + ip_int: int = Field(...) + latest: bool = Field(default=True) + owner: str = Field(...) + snapshots: list[Link["SnapshotDoc"]] = Field(default=[]) + source: str = Field(...) + time: datetime.datetime = Field(default_factory=datetime.datetime.utcnow) - @ip.setter - def ip(self, new_ip): - self._ip = new_ip - # Convert IP to an integer and store in ip_int - try: - self.ip_int = int(self._ip) - except ValueError: - raise ValueError(f"Invalid IP address: {new_ip}") + @model_validator(mode="before") + def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: + # ip may still be string if it was just set + values["ip_int"] = int(ipaddress.ip_address(values["ip"])) + return values - # Custom methods - def reset_latest_flag_by_owner(self, owner): - ScanDoc.objects(latest=True, owner=owner).update(latest=False) + class Config: + # Pydantic configuration + # Validate on assignment so ip_int is recalculated as ip is set + validate_assignment = True - def reset_latest_flag_by_ip(self, ips): - ip_ints = [int(x) for x in ips] if isinstance(ips, Iterable) else [int(ips)] - ScanDoc.objects(latest=True, ip_int__in=ip_ints).update(latest=False) + class Settings: + # Beanie settings + name = "scandocs" + indexes = [ + IndexModel( + [("latest", ASCENDING), ("ip_int", ASCENDING)], name="latest_ip" + ), + IndexModel([("time", ASCENDING), ("owner", ASCENDING)], name="time_owner"), + IndexModel([("int_ip", ASCENDING)], name="int_ip"), + IndexModel([("snapshots", ASCENDING)], name="snapshots", sparse=True), + ] - def tag_latest(self, owners, snapshot_oid): - ScanDoc.objects(latest=True, owner__in=owners).update( - push__snapshots=snapshot_oid + @classmethod + async def reset_latest_flag_by_owner(cls, owner): + await cls.find(cls.latest == True, cls.owner == owner).update( + Set({cls.latest: False}) ) - def tag_matching(self, existing_snapshot_oids, new_snapshot_oid): - ScanDoc.objects(snapshots__in=existing_snapshot_oids).update( - push__snapshots=new_snapshot_oid + @classmethod + async def reset_latest_flag_by_ip(cls, ips): + ip_ints = ( + [int(ipaddress.ip_address(x)) for x in ips] + if isinstance(ips, Iterable) + else [int(ipaddress.ip_address(ips))] ) - - def tag_timespan(self, owner, snapshot_oid, start_time, end_time): - ScanDoc.objects(time__gte=start_time, time__lte=end_time, owner=owner).update( - push__snapshots=snapshot_oid + await cls.find(cls.latest == True, cls.ip_int.in_(ip_ints)).update( + Set({cls.latest: False}) ) - def remove_tag(self, snapshot_oid): - ScanDoc.objects(snapshots=snapshot_oid).update(pull__snapshots=snapshot_oid) + @classmethod + async def tag_latest(cls, owners, snapshot_oid): + await cls.find(cls.latest == True, cls.owner.in_(owners)).update( + Push({cls.snapshots: snapshot_oid}) + ) diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index 72ba359..87613ca 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -1,5 +1,5 @@ # Third-Party Libraries -from mongoengine import Document +from beanie import Document class SnapshotDoc(Document): From 08920ddad9f8cab4a603c7ec51bf45c66d0357ce Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:46:02 -0500 Subject: [PATCH 014/139] Add cyhy_db package and initialize_db function --- src/cyhy_db/__init__.py | 3 +++ src/cyhy_db/db.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+) create mode 100644 src/cyhy_db/__init__.py create mode 100644 src/cyhy_db/db.py diff --git a/src/cyhy_db/__init__.py b/src/cyhy_db/__init__.py new file mode 100644 index 0000000..660c2b2 --- /dev/null +++ b/src/cyhy_db/__init__.py @@ -0,0 +1,3 @@ +from .db import initialize_db + +__all__ = ["initialize_db"] diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py new file mode 100644 index 0000000..9194655 --- /dev/null +++ b/src/cyhy_db/db.py @@ -0,0 +1,18 @@ +from beanie import init_beanie +from motor.motor_asyncio import AsyncIOMotorClient +from .models.cve import CVE +from .models.scan_doc import ScanDoc +from .models.snapshot_doc import SnapshotDoc + +ALL_MODELS = [CVE, ScanDoc, SnapshotDoc] + + +async def initialize_db(db_uri: str, db_name: str) -> None: + try: + client = AsyncIOMotorClient(db_uri) + db = client[db_name] + await init_beanie(database=db, document_models=ALL_MODELS) + return db + except Exception as e: + print(f"Failed to initialize database with error: {e}") + raise From 69117b44a8d5b83002ac90fab53e415c0b0cbc1d Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:46:45 -0500 Subject: [PATCH 015/139] Update conftest.py with Docker and mongo express fixtures --- tests/conftest.py | 57 +++++++++++++++++++++++++++++++++-------------- 1 file changed, 40 insertions(+), 17 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 85b69c2..561d280 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,16 +3,29 @@ https://docs.pytest.org/en/latest/writing_plugins.html#conftest-py-plugins """ +# Standard Python Libraries +import asyncio +import os +import time + # Third-Party Libraries +from beanie import Document, init_beanie import docker -from mongoengine import connect, disconnect +from motor.core import AgnosticClient +from motor.motor_asyncio import AsyncIOMotorClient import pytest +from pytest_asyncio import is_async_test + +# cisagov Libraries +from cyhy_db import initialize_db MONGO_INITDB_ROOT_USERNAME = os.environ.get("MONGO_INITDB_ROOT_USERNAME", "mongoadmin") MONGO_INITDB_ROOT_PASSWORD = os.environ.get("MONGO_INITDB_ROOT_PASSWORD", "secret") DATABASE_NAME = os.environ.get("DATABASE_NAME", "test") +MONGO_EXPRESS_PORT = os.environ.get("MONGO_EXPRESS_PORT", 8081) -docker_client = docker.from_env() +# Set the default event loop policy to be compatible with asyncio +AgnosticClient.get_io_loop = asyncio.get_running_loop @pytest.fixture(autouse=True) @@ -34,7 +47,13 @@ def group_github_log_lines(request): @pytest.fixture(scope="session") -def mongodb_container(mongo_image_tag): +def docker_client(): + """Fixture for the Docker client.""" + yield docker.from_env() + + +@pytest.fixture(scope="session") +def mongodb_container(docker_client, mongo_image_tag): """Fixture for the MongoDB test container.""" container = docker_client.containers.run( mongo_image_tag, @@ -74,8 +93,8 @@ def mongodb_container(mongo_image_tag): container.remove(force=True) -@pytest.fixture(scope="session") -def mongo_express_container(mongodb_container, request): +@pytest.fixture(autouse=True, scope="session") +def mongo_express_container(docker_client, request): if not request.config.getoption("--mongo-express"): yield None return @@ -96,7 +115,10 @@ def mongo_express_container(mongodb_container, request): def fin(): if request.config.getoption("--mongo-express"): - input("\n\nPress Enter to stop Mongo Express and MongoDB containers...") + print( + f"\n\nMongo Express is running at http://admin:pass@localhost:{MONGO_EXPRESS_PORT}" + ) + input("Press Enter to stop Mongo Express and MongoDB containers...") mongo_express_container.stop() mongo_express_container.remove(force=True) @@ -105,25 +127,26 @@ def fin(): @pytest.fixture(scope="session") -def mongo_uri(mongodb_container): - """Fixture for the MongoDB URI.""" +def db_uri(mongodb_container): + """Fixture for the database URI.""" mongo_port = mongodb_container.attrs["NetworkSettings"]["Ports"]["27017/tcp"][0][ "HostPort" ] - mongo_uri = f"mongodb://{MONGO_INITDB_ROOT_USERNAME}:{MONGO_INITDB_ROOT_PASSWORD}@localhost:{mongo_port}" - yield mongo_uri + uri = f"mongodb://{MONGO_INITDB_ROOT_USERNAME}:{MONGO_INITDB_ROOT_PASSWORD}@localhost:{mongo_port}" + yield uri @pytest.fixture(scope="session") -def mongodb_engine(mongo_uri, mongo_express_container): - """Fixture for the MongoDB engine.""" - connect( - host=mongo_uri, alias="default", db=DATABASE_NAME, uuidRepresentation="standard" - ) +def db_name(mongodb_container): + """Fixture for the database name.""" + yield DATABASE_NAME - yield - disconnect() +@pytest.fixture(autouse=True, scope="session") +async def db_client(db_uri): + """Fixture for client init.""" + print(f"Connecting to {db_uri}") + await initialize_db(db_uri, DATABASE_NAME) def pytest_addoption(parser): From 4d620e6c1046b1e96433c0bac472b527ec41a434 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:47:43 -0500 Subject: [PATCH 016/139] Add test for database connection using motor and beanie libraries --- tests/test_connection.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 tests/test_connection.py diff --git a/tests/test_connection.py b/tests/test_connection.py new file mode 100644 index 0000000..9771c5b --- /dev/null +++ b/tests/test_connection.py @@ -0,0 +1,21 @@ +"""Test database connection.""" + +# Third-Party Libraries +from motor.motor_asyncio import AsyncIOMotorClient + +from cyhy_db.models import CVE +from cyhy_db import initialize_db + + +async def test_connection_motor(db_uri, db_name): + client = AsyncIOMotorClient(db_uri) + db = client[db_name] + server_info = await db.command("ping") + assert server_info["ok"] == 1.0, "Direct database ping failed" + + +async def test_connection_beanie(): + # Attempt to find a document in the empty CVE collection + # await initialize_db(db_uri, db_name) # Manually initialize for testing + result = await CVE.get("CVE-2024-DOES-NOT-EXIST") + assert result is None, "Expected no document to be found" From c7959a9b787605384353522e2105e61fc7a43606 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:47:58 -0500 Subject: [PATCH 017/139] Refactor imports and update test cases in test_cve.py --- tests/test_cve.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tests/test_cve.py b/tests/test_cve.py index 51ba990..69926ad 100644 --- a/tests/test_cve.py +++ b/tests/test_cve.py @@ -1,11 +1,11 @@ """Test CVE model functionality.""" # Third-Party Libraries -from mongoengine import ValidationError +from pydantic import ValidationError import pytest # cisagov Libraries -from cyhy_db.models.cve import CVE +from cyhy_db.models import CVE severity_params = [ ("2.0", 10, 4), @@ -26,7 +26,7 @@ @pytest.mark.parametrize("version, score, expected_severity", severity_params) def test_calculate_severity(version, score, expected_severity): """Test that the severity is calculated correctly.""" - cve = CVE(cvss_version=version, cvss_score=score, id="test-cve") + cve = CVE(id="CVE-2024-0128", cvss_version=version, cvss_score=score) cve.calculate_severity() assert ( cve.severity == expected_severity @@ -36,16 +36,15 @@ def test_calculate_severity(version, score, expected_severity): @pytest.mark.parametrize("bad_score", [-1.0, 11.0]) def test_invalid_cvss_score(bad_score): """Test that an invalid CVSS score raises a ValueError.""" - cve = CVE(cvss_version="3.1", cvss_score=bad_score, id="test-cve") with pytest.raises(ValidationError): - cve.validate() # Explicitly call validate to trigger validation + CVE(cvss_version="3.1", cvss_score=bad_score, id="test-cve") -def test_save(mongodb_engine): +async def test_save(): """Test that the severity is calculated correctly on save.""" cve = CVE(cvss_version="3.1", cvss_score=9.0, id="test-cve") - cve.save() # Saving the object - saved_cve = CVE.objects(id="test-cve").first() # Retrieving the object + await cve.save() # Saving the object + saved_cve = await CVE.get("test-cve") # Retrieving the object assert saved_cve is not None, "CVE not saved correctly" assert saved_cve.severity == 4, "Severity not calculated correctly on save" From d551473d36be89d49fae7e0bda1b19df1e24d52a Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:48:07 -0500 Subject: [PATCH 018/139] Remove test_ip_address.py --- tests/test_ip_address.py | 60 ---------------------------------------- 1 file changed, 60 deletions(-) delete mode 100644 tests/test_ip_address.py diff --git a/tests/test_ip_address.py b/tests/test_ip_address.py deleted file mode 100644 index 893b381..0000000 --- a/tests/test_ip_address.py +++ /dev/null @@ -1,60 +0,0 @@ -# Standard Python Libraries -import ipaddress - -# Third-Party Libraries -from mongoengine import Document, ValidationError -import pytest - -# cisagov Libraries -from cyhy_db.models.ip_address import IPAddressField - - -class HasIpDocument(Document): - ip = IPAddressField() - - -def test_ip_address_type(): - doc = HasIpDocument(ip="1.2.3.4") - assert isinstance(doc.ip, (ipaddress.IPv4Address, ipaddress.IPv6Address)) - - -@pytest.mark.parametrize( - "valid_ip", - [ - "192.168.1.1", - "10.0.0.1", - "255.255.255.255", - ], -) -def test_valid_ip_address_field(valid_ip): - try: - HasIpDocument(ip=valid_ip).validate() - except ValidationError: - pytest.fail(f"Valid IP address {valid_ip} raised ValidationError") - - -@pytest.mark.parametrize( - "invalid_ip", - [ - "256.256.256.256", - "123.456.789.0", - "abc.def.ghi.jkl", - ], -) -def test_invalid_ip_address_field(invalid_ip): - with pytest.raises(ValidationError): - HasIpDocument(ip=invalid_ip).validate() - - -def test_save_ip_address(mongodb_engine): - test_document = HasIpDocument(ip=ipaddress.IPv4Address("192.168.1.1")) - test_document.save() - assert test_document.id is not None, "TestDocument instance was not saved correctly" - - -def test_retrieve_ip_address(mongodb_engine): - retrieved_doc = HasIpDocument.objects().first() - assert retrieved_doc is not None, "TestDocument instance was not retrieved" - assert retrieved_doc.ip == ipaddress.IPv4Address( - "192.168.1.1" - ), "Retrieved IP address does not match the saved IP address" From 577ee7d890429207bcc4ad275a5bb5be275eb6de Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:48:15 -0500 Subject: [PATCH 019/139] Add tests for ScanDoc model functionality --- tests/test_scandoc.py | 86 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 86 insertions(+) create mode 100644 tests/test_scandoc.py diff --git a/tests/test_scandoc.py b/tests/test_scandoc.py new file mode 100644 index 0000000..6ef7cb2 --- /dev/null +++ b/tests/test_scandoc.py @@ -0,0 +1,86 @@ +"""Test ScanDoc model functionality.""" + +import ipaddress + +# Third-Party Libraries +import pytest +from pydantic import ValidationError + +# cisagov Libraries +from cyhy_db.models import ScanDoc + + +VALID_IP_1_STR = "0.0.0.1" +VALID_IP_2_STR = "0.0.0.2" +VALID_IP_1_INT = int(ipaddress.ip_address(VALID_IP_1_STR)) +VALID_IP_2_INT = int(ipaddress.ip_address(VALID_IP_2_STR)) + + +def test_ip_int_init(): + # Create a ScanDoc object + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner="YOUR_MOM", + source="nmap", + ) + + assert scan_doc.ip_int == int( + ipaddress.ip_address(VALID_IP_1_STR) + ), "IP address integer was not calculated correctly on init" + + +def test_ip_int_change(): + # Create a ScanDoc object + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner="YOUR_MOM", + source="nmap", + ) + + scan_doc.ip = ipaddress.ip_address(VALID_IP_2_STR) + + assert scan_doc.ip_int == int( + ipaddress.ip_address(VALID_IP_2_STR) + ), "IP address integer was not calculated correctly on change" + + +def test_ip_string_set(): + scan_doc = ScanDoc( + ip=VALID_IP_1_STR, + owner="YOUR_MOM", + source="nmap", + ) + + assert type(scan_doc.ip) == ipaddress.IPv4Address, "IP address was not converted" + assert scan_doc.ip_int == VALID_IP_1_INT, "IP address integer was not calculated" + + +async def test_ip_address_field_fetch(): + # Create a ScanDoc object + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner="YOUR_MOM", + source="nmap", + ) + + # Save the ScanDoc object to the database + await scan_doc.save() + + # Retrieve the ScanDoc object from the database + retrieved_doc = await ScanDoc.get(scan_doc.id) + + # Assert that the retrieved IP address is equal to the one we saved + assert retrieved_doc.ip == ipaddress.ip_address( + VALID_IP_1_STR + ), "IP address does not match" + + assert retrieved_doc.ip_int == VALID_IP_1_INT, "IP address integer does not match" + + +def test_invalid_ip_address(): + with pytest.raises(ValidationError): + ScanDoc( + ip="999.999.999.999", # This should be invalid + owner="owner_example", + source="source_example", + ) From f2e99a3bc4c47fca51e63f1818353f956851e93b Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:48:24 -0500 Subject: [PATCH 020/139] Add asyncio_mode to pytest.ini --- pytest.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/pytest.ini b/pytest.ini index ed958e0..189f9fc 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,2 +1,3 @@ [pytest] addopts = -v -ra --cov +asyncio_mode = auto From 89e8303989b3c5f105a351cb9ac4b1fc722fe63c Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 7 Feb 2024 11:48:55 -0500 Subject: [PATCH 021/139] Switch to beanie ORM as a dependency --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index e2de6f5..cf5c127 100644 --- a/setup.py +++ b/setup.py @@ -91,8 +91,8 @@ def get_version(version_file): py_modules=[splitext(basename(path))[0] for path in glob("src/*.py")], include_package_data=True, install_requires=[ + "beanie == 1.25.0", "docopt == 0.6.2", - "mongoengine == 0.27.0", "schema == 0.7.5", "setuptools >= 69.0.3", ], From 2cf458712a1f7edf5211b0dba4421fc3ac1ef6ba Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 12 Feb 2024 13:12:16 -0500 Subject: [PATCH 022/139] Add enum classes for various types and statuses --- src/cyhy_db/models/enum.py | 89 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 src/cyhy_db/models/enum.py diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py new file mode 100644 index 0000000..cbd49bd --- /dev/null +++ b/src/cyhy_db/models/enum.py @@ -0,0 +1,89 @@ +from enum import Enum + + +class AgencyType(Enum): + FEDERAL = "FEDERAL" + LOCAL = "LOCAL" + PRIVATE = "PRIVATE" + STATE = "STATE" + TERRITORIAL = "TERRITORIAL" + TRIBAL = "TRIBAL" + + +class ControlAction(Enum): + PAUSE = "PAUSE" + STOP = "STOP" + + +class ControlTarget(Enum): + COMMANDER = "COMMANDER" + + +class CVSSVersion(Enum): + V2 = "2.0" + V3 = "3.0" + V3_1 = "3.1" + + +class DayOfWeek(Enum): + MONDAY = "MONDAY" + TUESDAY = "TUESDAY" + WEDNESDAY = "WEDNESDAY" + THURSDAY = "THURSDAY" + FRIDAY = "FRIDAY" + SATURDAY = "SATURDAY" + SUNDAY = "SUNDAY" + + +class PocType(Enum): + DISTRO = "DISTRO" + TECHNICAL = "TECHNICAL" + + +class ReportPeriod(Enum): + MONTHLY = "MONTHLY" + QUARTERLY = "QUARTERLY" + WEEKLY = "WEEKLY" + + +class ReportType(Enum): + BOD = "BOD" + CYBEX = "CYBEX" + CYHY = "CYHY" + CYHY_THIRD_PARTY = "CYHY_THIRD_PARTY" + DNSSEC = "DNSSEC" + PHISHING = "PHISHING" + + +class ScanType(Enum): + CYHY = "CYHY" + DNSSEC = "DNSSEC" + PHISHING = "PHISHING" + + +class Scheduler(Enum): + PERSISTENT1 = "PERSISTENT1" + + +class Stage(Enum): + BASESCAN = "BASESCAN" + NETSCAN1 = "NETSCAN1" + NETSCAN2 = "NETSCAN2" + PORTSCAN = "PORTSCAN" + VULNSCAN = "VULNSCAN" + + +class Status(Enum): + DONE = "DONE" + READY = "READY" + RUNNING = "RUNNING" + WAITING = "WAITING" + + +class TicketEvent(Enum): + CHANGED = "CHANGED" + CLOSED = "CLOSED" + OPENED = "OPENED" + REOPENED = "REOPENED" + UNVERIFIED = "UNVERIFIED" + VERIFIED = "VERIFIED" From 43da94be617b2336a3f80c70915e70c09e0f8d96 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 12 Feb 2024 13:12:56 -0500 Subject: [PATCH 023/139] Add RequestDoc model and test --- src/cyhy_db/models/__init__.py | 3 +- src/cyhy_db/models/request_doc.py | 94 +++++++++++++++++++++++++++++++ tests/test_request_doc.py | 34 +++++++++++ 3 files changed, 130 insertions(+), 1 deletion(-) create mode 100644 src/cyhy_db/models/request_doc.py create mode 100644 tests/test_request_doc.py diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py index 87d45e9..c4608ef 100644 --- a/src/cyhy_db/models/__init__.py +++ b/src/cyhy_db/models/__init__.py @@ -1,5 +1,6 @@ from .cve import CVE from .scan_doc import ScanDoc from .snapshot_doc import SnapshotDoc +from .request_doc import RequestDoc -__all__ = ["CVE", "ScanDoc", "SnapshotDoc"] +__all__ = ["CVE", "RequestDoc", "ScanDoc", "SnapshotDoc"] diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py new file mode 100644 index 0000000..23a6a77 --- /dev/null +++ b/src/cyhy_db/models/request_doc.py @@ -0,0 +1,94 @@ +from pydantic import BaseModel, Field, field_validator, EmailStr +from datetime import datetime +from typing import List, Optional +from beanie import before_event, Document, Insert, Link, Replace, ValidateOnSave +from ipaddress import IPv4Network +import re + +from .enum import ( + AgencyType, + DayOfWeek, + PocType, + ReportPeriod, + ReportType, + ScanType, + Scheduler, + Stage, +) + +BOGUS_ID = "bogus_id_replace_me" + + +class Contact(BaseModel): + email: EmailStr + name: str + phone: str + type: PocType + + +class Location(BaseModel): + country_name: str + country: str + county_fips: str + county: str + gnis_id: int + name: str + state_fips: str + state_name: str + state: str + + +class Agency(BaseModel): + name: str + acronym: str + type: Optional[AgencyType] = Field(default=None) + contacts: List[Contact] = Field(default=[]) + location: Optional[Location] = Field(default=None) + + +class ScanLimit(BaseModel): + scan_type: ScanType = Field(..., alias="scanType") + concurrent: int = Field(ge=0) + + +class Window(BaseModel): + day: DayOfWeek = Field(default=DayOfWeek.SUNDAY) + duration: int = Field(default=168, ge=0, le=168) + start: str = Field(default="00:00:00") + + @field_validator("start") + def validate_start(cls, v): + # Validate that the start time is in the format HH:MM:SS + if not re.match(r"^\d{2}:\d{2}:\d{2}$", v): + raise ValueError("Start time must be in the format HH:MM:SS") + return v + + +class RequestDoc(Document): + id: str = Field(default=BOGUS_ID) + agency: Agency + children: List[Link["RequestDoc"]] = Field(default=[]) + enrolled: datetime = Field(default_factory=datetime.utcnow) + init_stage: Stage = Field(default=Stage.NETSCAN1) + key: Optional[str] = Field(default=None) + networks: List[IPv4Network] = Field(default=[]) + period_start: datetime = Field(default_factory=datetime.utcnow) + report_period: ReportPeriod = Field(default=ReportPeriod.WEEKLY) + report_types: List[ReportType] = Field(default=[]) + retired: bool = False + scan_limits: List[ScanLimit] = Field(default=[]) + scan_types: List[ScanType] = Field(default=[]) + scheduler: Scheduler = Field(default=Scheduler.PERSISTENT1) + stakeholder: bool = False + windows: List[Window] = Field(default=[Window()]) + + @before_event(Insert, Replace, ValidateOnSave) + async def set_id_to_acronym(self): + # Set the id to the agency acronym if it is the default value + if self.id == BOGUS_ID: + self.id = self.agency.acronym + + class Settings: + # Beanie settings + name = "requests" + indexes = [] diff --git a/tests/test_request_doc.py b/tests/test_request_doc.py new file mode 100644 index 0000000..0df0466 --- /dev/null +++ b/tests/test_request_doc.py @@ -0,0 +1,34 @@ +"""Test RequestDoc model functionality.""" + +import ipaddress + +# Third-Party Libraries +import pytest +from pydantic import ValidationError +from hypothesis import given, strategies as st + +# cisagov Libraries +from cyhy_db.models import RequestDoc +from cyhy_db.models.request_doc import Agency + + +async def test_init(): + # Create a RequestDoc object + + request_doc = RequestDoc( + agency=Agency( + name="Cybersecurity and Infrastructure Security Agency", acronym="CISA" + ) + ) + + await request_doc.save() + + # Verify that the id was set to the acronym + assert ( + request_doc.id == request_doc.agency.acronym + ), "id was not correctly set to agency acronym" + + +# @given(st.builds(RequestDoc)) +# def test_dump_model(instance): +# print(instance) From 4162bdb748faece3af3704610937226e214b2ab4 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 12 Feb 2024 13:13:24 -0500 Subject: [PATCH 024/139] Refactor CVE model and add model validator for calculating severity --- src/cyhy_db/models/cve.py | 53 +++++++++++++++++++++++---------------- tests/test_cve.py | 1 - 2 files changed, 31 insertions(+), 23 deletions(-) diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 54828c8..c421dd5 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -1,34 +1,43 @@ # Third-Party Libraries from beanie import Document, Indexed, ValidateOnSave, before_event -from pydantic import Field +from pydantic import Field, model_validator +from .enum import CVSSVersion +from typing import Any, Dict class CVE(Document): id: str = Indexed(primary_field=True) # CVE ID cvss_score: float = Field(ge=0.0, le=10.0) - cvss_version: str = Field(enum=["2.0", "3.0", "3.1"]) + cvss_version: CVSSVersion = Field(default=CVSSVersion.V3_1) severity: int = Field(ge=1, le=4, default=1) - class Settings: - name = "cves" - - @before_event(ValidateOnSave) - def calculate_severity(self): - if self.cvss_version == "2.0": - if self.cvss_score == 10: - self.severity = 4 - elif self.cvss_score >= 7.0: - self.severity = 3 - elif self.cvss_score >= 4.0: - self.severity = 2 + @model_validator(mode="before") + def calculate_severity(cls, values: Dict[str, Any]) -> Dict[str, Any]: + if values["cvss_version"] == "2.0": + if values["cvss_score"] == 10: + values["severity"] = 4 + elif values["cvss_score"] >= 7.0: + values["severity"] = 3 + elif values["cvss_score"] >= 4.0: + values["severity"] = 2 else: - self.severity = 1 + values["severity"] = 1 else: # CVSS versions 3.0 or 3.1 - if self.cvss_score >= 9.0: - self.severity = 4 - elif self.cvss_score >= 7.0: - self.severity = 3 - elif self.cvss_score >= 4.0: - self.severity = 2 + if values["cvss_score"] >= 9.0: + values["severity"] = 4 + elif values["cvss_score"] >= 7.0: + values["severity"] = 3 + elif values["cvss_score"] >= 4.0: + values["severity"] = 2 else: - self.severity = 1 + values["severity"] = 1 + return values + + class Config: + # Pydantic configuration + # Validate on assignment so ip_int is recalculated as ip is set + validate_assignment = True + + class Settings: + # Beanie settings + name = "cves" diff --git a/tests/test_cve.py b/tests/test_cve.py index 69926ad..d9f999d 100644 --- a/tests/test_cve.py +++ b/tests/test_cve.py @@ -27,7 +27,6 @@ def test_calculate_severity(version, score, expected_severity): """Test that the severity is calculated correctly.""" cve = CVE(id="CVE-2024-0128", cvss_version=version, cvss_score=score) - cve.calculate_severity() assert ( cve.severity == expected_severity ), f"Failed for CVSS {version} with score {score}" From c45dd4848596b852bcd548bf6f3be3f40ab22261 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 12 Feb 2024 13:14:12 -0500 Subject: [PATCH 025/139] Rename ScanDoc test --- tests/{test_scandoc.py => test_scan_doc.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{test_scandoc.py => test_scan_doc.py} (100%) diff --git a/tests/test_scandoc.py b/tests/test_scan_doc.py similarity index 100% rename from tests/test_scandoc.py rename to tests/test_scan_doc.py From a26b5c27953c2fe38b1bdbf8cbe332ddbd21f791 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 12 Feb 2024 13:14:37 -0500 Subject: [PATCH 026/139] Add RequestDoc model to ALL_MODELS list --- src/cyhy_db/db.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index 9194655..867c292 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -1,10 +1,11 @@ from beanie import init_beanie from motor.motor_asyncio import AsyncIOMotorClient from .models.cve import CVE +from .models.request_doc import RequestDoc from .models.scan_doc import ScanDoc from .models.snapshot_doc import SnapshotDoc -ALL_MODELS = [CVE, ScanDoc, SnapshotDoc] +ALL_MODELS = [CVE, RequestDoc, ScanDoc, SnapshotDoc] async def initialize_db(db_uri: str, db_name: str) -> None: From 11b5bb22262873fc5dc371d408171a2d2dc4cb2b Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 12 Feb 2024 13:15:00 -0500 Subject: [PATCH 027/139] Add new dependencies to setup.py --- setup.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/setup.py b/setup.py index cf5c127..37f74f2 100644 --- a/setup.py +++ b/setup.py @@ -93,6 +93,7 @@ def get_version(version_file): install_requires=[ "beanie == 1.25.0", "docopt == 0.6.2", + "pydantic[email, hypothesis]", "schema == 0.7.5", "setuptools >= 69.0.3", ], @@ -108,8 +109,12 @@ def get_version(version_file): # to never grab the regression version. "coveralls != 1.11.0", "docker == 7.0.0", + "hypothesis", + "mimesis-factory", + "mimesis", "pre-commit", "pytest-cov", + "pytest-factoryboy", "pytest", ] }, From 1b1d518ca0ab3b83acd4621e2559f54cbdface5a Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 12 Feb 2024 13:49:53 -0500 Subject: [PATCH 028/139] Add test case to reset latest flag by owner --- tests/test_scan_doc.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index 6ef7cb2..9099253 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -84,3 +84,19 @@ def test_invalid_ip_address(): owner="owner_example", source="source_example", ) + + +async def test_reset_latest_flag_by_owner(): + # Create a ScanDoc object + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), owner="RESET_MY_LATEST", source="nmap" + ) + await scan_doc.save() + # Check that the latest flag is set to True + assert scan_doc.latest == True + # Reset the latest flag + await ScanDoc.reset_latest_flag_by_owner("RESET_MY_LATEST") + # Retrieve the ScanDoc object from the database + await scan_doc.sync() + # Check that the latest flag is set to False + assert scan_doc.latest == False From 2433f21b8e15382188f9394f5d50832bb1beaccf Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 13 Feb 2024 12:46:56 -0500 Subject: [PATCH 029/139] Sort imports --- src/cyhy_db/models/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py index c4608ef..d42b750 100644 --- a/src/cyhy_db/models/__init__.py +++ b/src/cyhy_db/models/__init__.py @@ -1,6 +1,6 @@ from .cve import CVE +from .request_doc import RequestDoc from .scan_doc import ScanDoc from .snapshot_doc import SnapshotDoc -from .request_doc import RequestDoc __all__ = ["CVE", "RequestDoc", "ScanDoc", "SnapshotDoc"] From 7adf2d104eb2e9daa19d2a19c87e535ffc456096 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 13 Feb 2024 12:47:46 -0500 Subject: [PATCH 030/139] Refactor ScanDoc model and add new tests --- src/cyhy_db/models/scan_doc.py | 56 ++++++++++++++++++++++--------- tests/test_scan_doc.py | 60 ++++++++++++++++++++++++++++++++-- 2 files changed, 98 insertions(+), 18 deletions(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 6d3af58..79c4fc3 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -1,11 +1,13 @@ # Standard Python Libraries import datetime import ipaddress -from typing import Any, Dict, Iterable +from typing import Any, Dict, Iterable, List, Union # Third-Party Libraries from beanie import Document, Link -from beanie.operators import Push, Set +from beanie.operators import Push, In, Set +from bson import ObjectId +from bson.dbref import DBRef from pydantic import Field, model_validator from pymongo import ASCENDING, IndexModel @@ -15,7 +17,7 @@ class ScanDoc(Document): ip_int: int = Field(...) latest: bool = Field(default=True) owner: str = Field(...) - snapshots: list[Link["SnapshotDoc"]] = Field(default=[]) + snapshots: List[Link["SnapshotDoc"]] = Field(default=[]) source: str = Field(...) time: datetime.datetime = Field(default_factory=datetime.datetime.utcnow) @@ -43,24 +45,48 @@ class Settings: ] @classmethod - async def reset_latest_flag_by_owner(cls, owner): - await cls.find(cls.latest == True, cls.owner == owner).update( + async def reset_latest_flag_by_owner(cls, owner: str): + await cls.find(cls.latest == True, cls.owner == owner).update_many( Set({cls.latest: False}) ) @classmethod - async def reset_latest_flag_by_ip(cls, ips): - ip_ints = ( - [int(ipaddress.ip_address(x)) for x in ips] - if isinstance(ips, Iterable) - else [int(ipaddress.ip_address(ips))] - ) - await cls.find(cls.latest == True, cls.ip_int.in_(ip_ints)).update( + async def reset_latest_flag_by_ip( + cls, + ips: ( + int + | ipaddress.IPv4Address + | Iterable[int] + | Iterable[ipaddress.IPv4Address] + | Iterable[str] + | str + ), + ): + + if isinstance(ips, Iterable): + # TODO Figure out why coverage thinks this next line can exit early + ip_ints = [int(ipaddress.ip_address(x)) for x in ips] + else: + ip_ints = [int(ipaddress.ip_address(ips))] + + await cls.find(cls.latest == True, In(cls.ip_int, ip_ints)).update_many( Set({cls.latest: False}) ) @classmethod - async def tag_latest(cls, owners, snapshot_oid): - await cls.find(cls.latest == True, cls.owner.in_(owners)).update( - Push({cls.snapshots: snapshot_oid}) + async def tag_latest( + cls, owners: List[str], snapshot: Union["SnapshotDoc", ObjectId, str] + ): + from . import SnapshotDoc + + if isinstance(snapshot, SnapshotDoc): + ref = DBRef(SnapshotDoc.Settings.name, snapshot.id) + elif isinstance(snapshot, ObjectId): + ref = DBRef(SnapshotDoc.Settings.name, snapshot) + elif isinstance(snapshot, str): + ref = DBRef(SnapshotDoc.Settings.name, ObjectId(snapshot)) + else: + raise ValueError("Invalid snapshot type") + await cls.find(cls.latest == True, In(cls.owner, owners)).update_many( + Push({cls.snapshots: ref}) ) diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index 9099253..f34a28b 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -1,13 +1,14 @@ """Test ScanDoc model functionality.""" import ipaddress +import datetime # Third-Party Libraries import pytest from pydantic import ValidationError # cisagov Libraries -from cyhy_db.models import ScanDoc +from cyhy_db.models import ScanDoc, SnapshotDoc VALID_IP_1_STR = "0.0.0.1" @@ -88,15 +89,68 @@ def test_invalid_ip_address(): async def test_reset_latest_flag_by_owner(): # Create a ScanDoc object + OWNER = "RESET_BY_OWNER" scan_doc = ScanDoc( - ip=ipaddress.ip_address(VALID_IP_1_STR), owner="RESET_MY_LATEST", source="nmap" + ip=ipaddress.ip_address(VALID_IP_1_STR), owner=OWNER, source="nmap" ) await scan_doc.save() # Check that the latest flag is set to True assert scan_doc.latest == True # Reset the latest flag - await ScanDoc.reset_latest_flag_by_owner("RESET_MY_LATEST") + await ScanDoc.reset_latest_flag_by_owner(OWNER) # Retrieve the ScanDoc object from the database await scan_doc.sync() # Check that the latest flag is set to False assert scan_doc.latest == False + + +async def test_reset_latest_flag_by_ip(): + # Create a ScanDoc object + IP_TO_RESET_1 = ipaddress.ip_address("128.205.1.2") + IP_TO_RESET_2 = ipaddress.ip_address("128.205.1.3") + scan_doc_1 = ScanDoc(ip=IP_TO_RESET_1, owner="RESET_BY_IP", source="nmap") + scan_doc_2 = ScanDoc(ip=IP_TO_RESET_2, owner="RESET_BY_IP", source="nmap") + await scan_doc_1.save() + await scan_doc_2.save() + # Check that the latest flag is set to True + assert scan_doc_1.latest == True + # Reset the latest flag on single IP + await ScanDoc.reset_latest_flag_by_ip(IP_TO_RESET_1) + # Retrieve the ScanDoc object from the database + await scan_doc_1.sync() + # Check that the latest flag is set to False + assert scan_doc_1.latest == False + # Reset by both IPs + await ScanDoc.reset_latest_flag_by_ip([IP_TO_RESET_1, IP_TO_RESET_2]) + # Retrieve the ScanDoc object from the database + await scan_doc_2.sync() + # Check that the latest flag is set to False + assert scan_doc_2.latest == False + + +async def test_tag_latest(): + # Create a SnapshotDoc object + + owner = "TAG_LATEST" + snapshot_doc = SnapshotDoc( + owner=owner, + start_time=datetime.datetime.utcnow(), + end_time=datetime.datetime.utcnow(), + ) + await snapshot_doc.save() + # Create a ScanDoc object + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner=owner, + source="nmap", + ) + await scan_doc.save() + + # Tag the latest scan + await ScanDoc.tag_latest([owner], snapshot_doc) + + # Retrieve the ScanDoc object from the database + scan_doc = await ScanDoc.find_one(ScanDoc.id == scan_doc.id, fetch_links=True) + + # Check that the scan now has a snapshot + assert scan_doc.snapshots == [snapshot_doc], "Snapshot not added to scan" From 14a084846e6021d37ce01da7e3b300fafae531bd Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 13 Feb 2024 12:48:07 -0500 Subject: [PATCH 031/139] Add fields and indexes to SnapshotDoc model --- src/cyhy_db/models/snapshot_doc.py | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index 87613ca..3f1197f 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -1,6 +1,29 @@ # Third-Party Libraries from beanie import Document +from pydantic import Field +from pymongo import ASCENDING, IndexModel +from datetime import datetime class SnapshotDoc(Document): - pass + owner: str = Field(...) + start_time: datetime = Field(...) + end_time: datetime = Field(...) + + class Settings: + # Beanie settings + name = "snapshots" + indexes = [ + IndexModel( + [ + ("owner", ASCENDING), + ("start_time", ASCENDING), + ("end_time", ASCENDING), + ], + name="uniques", + unique=True, + ), + IndexModel( + [("latest", ASCENDING), ("owner", ASCENDING)], name="latest_owner" + ), + ] From 718f4d568173500e7347e2619796f83748c52b30 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 13 Feb 2024 12:49:55 -0500 Subject: [PATCH 032/139] Organize imports --- src/cyhy_db/db.py | 2 ++ src/cyhy_db/models/cve.py | 5 ++++- src/cyhy_db/models/enum.py | 1 + src/cyhy_db/models/request_doc.py | 9 ++++++--- src/cyhy_db/models/scan_doc.py | 3 +-- src/cyhy_db/models/snapshot_doc.py | 4 +++- tests/test_connection.py | 3 ++- tests/test_request_doc.py | 6 ++++-- tests/test_scan_doc.py | 6 +++--- 9 files changed, 26 insertions(+), 13 deletions(-) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index 867c292..9698924 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -1,5 +1,7 @@ +# Third-Party Libraries from beanie import init_beanie from motor.motor_asyncio import AsyncIOMotorClient + from .models.cve import CVE from .models.request_doc import RequestDoc from .models.scan_doc import ScanDoc diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index c421dd5..0c0489c 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -1,8 +1,11 @@ +# Standard Python Libraries +from typing import Any, Dict + # Third-Party Libraries from beanie import Document, Indexed, ValidateOnSave, before_event from pydantic import Field, model_validator + from .enum import CVSSVersion -from typing import Any, Dict class CVE(Document): diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index cbd49bd..681a4c9 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -1,3 +1,4 @@ +# Standard Python Libraries from enum import Enum diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 23a6a77..277085b 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -1,9 +1,12 @@ -from pydantic import BaseModel, Field, field_validator, EmailStr +# Standard Python Libraries from datetime import datetime -from typing import List, Optional -from beanie import before_event, Document, Insert, Link, Replace, ValidateOnSave from ipaddress import IPv4Network import re +from typing import List, Optional + +# Third-Party Libraries +from beanie import Document, Insert, Link, Replace, ValidateOnSave, before_event +from pydantic import BaseModel, EmailStr, Field, field_validator from .enum import ( AgencyType, diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 79c4fc3..a10bbfb 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -5,7 +5,7 @@ # Third-Party Libraries from beanie import Document, Link -from beanie.operators import Push, In, Set +from beanie.operators import In, Push, Set from bson import ObjectId from bson.dbref import DBRef from pydantic import Field, model_validator @@ -62,7 +62,6 @@ async def reset_latest_flag_by_ip( | str ), ): - if isinstance(ips, Iterable): # TODO Figure out why coverage thinks this next line can exit early ip_ints = [int(ipaddress.ip_address(x)) for x in ips] diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index 3f1197f..5691999 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -1,8 +1,10 @@ +# Standard Python Libraries +from datetime import datetime + # Third-Party Libraries from beanie import Document from pydantic import Field from pymongo import ASCENDING, IndexModel -from datetime import datetime class SnapshotDoc(Document): diff --git a/tests/test_connection.py b/tests/test_connection.py index 9771c5b..a26f9fb 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -3,8 +3,9 @@ # Third-Party Libraries from motor.motor_asyncio import AsyncIOMotorClient -from cyhy_db.models import CVE +# cisagov Libraries from cyhy_db import initialize_db +from cyhy_db.models import CVE async def test_connection_motor(db_uri, db_name): diff --git a/tests/test_request_doc.py b/tests/test_request_doc.py index 0df0466..d4e627d 100644 --- a/tests/test_request_doc.py +++ b/tests/test_request_doc.py @@ -1,11 +1,13 @@ """Test RequestDoc model functionality.""" +# Standard Python Libraries import ipaddress # Third-Party Libraries -import pytest +from hypothesis import given +from hypothesis import strategies as st from pydantic import ValidationError -from hypothesis import given, strategies as st +import pytest # cisagov Libraries from cyhy_db.models import RequestDoc diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index f34a28b..1ca4b46 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -1,16 +1,16 @@ """Test ScanDoc model functionality.""" -import ipaddress +# Standard Python Libraries import datetime +import ipaddress # Third-Party Libraries -import pytest from pydantic import ValidationError +import pytest # cisagov Libraries from cyhy_db.models import ScanDoc, SnapshotDoc - VALID_IP_1_STR = "0.0.0.1" VALID_IP_2_STR = "0.0.0.2" VALID_IP_1_INT = int(ipaddress.ip_address(VALID_IP_1_STR)) From 74de4e6b09730eadd6c7339e2890df19bb32d374 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 13 Feb 2024 17:24:18 -0500 Subject: [PATCH 033/139] Add HostDoc model and test cases --- src/cyhy_db/models/__init__.py | 3 +- src/cyhy_db/models/host_doc.py | 90 ++++++++++++++++++++++++++++++++++ tests/test_host_doc.py | 33 +++++++++++++ 3 files changed, 125 insertions(+), 1 deletion(-) create mode 100644 src/cyhy_db/models/host_doc.py create mode 100644 tests/test_host_doc.py diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py index d42b750..723a1de 100644 --- a/src/cyhy_db/models/__init__.py +++ b/src/cyhy_db/models/__init__.py @@ -1,6 +1,7 @@ from .cve import CVE +from .host_doc import HostDoc from .request_doc import RequestDoc from .scan_doc import ScanDoc from .snapshot_doc import SnapshotDoc -__all__ = ["CVE", "RequestDoc", "ScanDoc", "SnapshotDoc"] +__all__ = ["CVE", "HostDoc", "RequestDoc", "ScanDoc", "SnapshotDoc"] diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py new file mode 100644 index 0000000..4266008 --- /dev/null +++ b/src/cyhy_db/models/host_doc.py @@ -0,0 +1,90 @@ +from beanie import Document, before_event, Indexed, Insert, Replace, ValidateOnSave +from datetime import datetime +from pydantic import BaseModel, Field, model_validator +from pymongo import ASCENDING, IndexModel +from typing import Any, Dict, Optional, Tuple +import random +from .enum import Stage, Status +from ipaddress import IPv4Address, ip_address + + +class State(BaseModel): + reason: str + up: bool + + +class HostDoc(Document): + id: int = Field(...) # IP address as an integer + ip: IPv4Address = Field(...) + owner: str = Field(...) + last_change: datetime = Field(default_factory=datetime.utcnow) + next_scan: Optional[datetime] = Field(default=None) + state: State = Field(default_factory=lambda: State(reason="new", up=False)) + stage: Stage = Field(default=Stage.NETSCAN1) + status: Status = Field(default=Status.WAITING) + loc: Optional[Tuple[float, float]] = Field(default=None) + priority: int = Field(default=0) + r: float = Field(default_factory=random.random) + latest_scan: Dict[Stage, datetime] = Field(default_factory=dict) + + @model_validator(mode="before") + def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: + # ip may still be string if it was just set + values["id"] = int(ip_address(values["ip"])) + print(values) + return values + + @before_event(Insert, Replace, ValidateOnSave) + async def before_save(self): + self.last_change = datetime.utcnow() + + class Settings: + name = "hosts" + indexes = [ + IndexModel( + [ + ("status", ASCENDING), + ("stage", ASCENDING), + ("owner", ASCENDING), + ("priority", ASCENDING), + ("r", ASCENDING), + ], + name="claim", + ), + IndexModel( + [ + ("ip", ASCENDING), + ], + name="ip", + ), + IndexModel( + [ + ("state.up", ASCENDING), + ("owner", ASCENDING), + ], + name="up", + ), + IndexModel( + [ + ("next_scan", ASCENDING), + ("state.up", ASCENDING), + ("status", ASCENDING), + ], + sparse=True, + name="next_scan", + ), + IndexModel( + [ + ("owner", ASCENDING), + ], + name="owner", + ), + IndexModel( + [ + ("owner", ASCENDING), + ("state.up", ASCENDING), + ("latest_scan.VULNSCAN", ASCENDING), + ], + name="latest_scan_done", + ), + ] diff --git a/tests/test_host_doc.py b/tests/test_host_doc.py new file mode 100644 index 0000000..a1e0a57 --- /dev/null +++ b/tests/test_host_doc.py @@ -0,0 +1,33 @@ +from cyhy_db.models import HostDoc +from ipaddress import ip_address + +VALID_IP_1_STR = "0.0.0.1" +VALID_IP_2_STR = "0.0.0.2" +VALID_IP_1_INT = int(ip_address(VALID_IP_1_STR)) +VALID_IP_2_INT = int(ip_address(VALID_IP_2_STR)) + + +def test_host_doc_init(): + # Create a HostDoc object + host_doc = HostDoc( + ip=ip_address(VALID_IP_1_STR), + owner="YOUR_MOM", + ) + + +async def test_save(): + # Create a HostDoc object + host_doc = HostDoc( + ip=ip_address(VALID_IP_1_STR), + owner="YOUR_MOM", + ) + # Save the HostDoc object to the database + await host_doc.save() + assert host_doc.id == VALID_IP_1_INT + + +async def test_find(): + # Find a HostDoc object by its IP address + host_doc = await HostDoc.find_one(HostDoc.id == VALID_IP_1_INT) + print(host_doc) + assert host_doc.ip == ip_address(VALID_IP_1_STR) From b28e54a0dc0d80a1a62c5fd8ec608dcc4214fd0e Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 13 Feb 2024 17:24:38 -0500 Subject: [PATCH 034/139] Update imports and type annotations in scan_doc.py --- src/cyhy_db/models/scan_doc.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index a10bbfb..8421f8f 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -1,6 +1,6 @@ # Standard Python Libraries -import datetime -import ipaddress +from datetime import datetime +from ipaddress import IPv4Address, ip_address from typing import Any, Dict, Iterable, List, Union # Third-Party Libraries @@ -13,18 +13,18 @@ class ScanDoc(Document): - ip: ipaddress.IPv4Address = Field(...) + ip: IPv4Address = Field(...) ip_int: int = Field(...) latest: bool = Field(default=True) owner: str = Field(...) snapshots: List[Link["SnapshotDoc"]] = Field(default=[]) source: str = Field(...) - time: datetime.datetime = Field(default_factory=datetime.datetime.utcnow) + time: datetime = Field(default_factory=datetime.utcnow) @model_validator(mode="before") def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: # ip may still be string if it was just set - values["ip_int"] = int(ipaddress.ip_address(values["ip"])) + values["ip_int"] = int(ip_address(values["ip"])) return values class Config: @@ -55,18 +55,18 @@ async def reset_latest_flag_by_ip( cls, ips: ( int - | ipaddress.IPv4Address + | IPv4Address | Iterable[int] - | Iterable[ipaddress.IPv4Address] + | Iterable[IPv4Address] | Iterable[str] | str ), ): if isinstance(ips, Iterable): # TODO Figure out why coverage thinks this next line can exit early - ip_ints = [int(ipaddress.ip_address(x)) for x in ips] + ip_ints = [int(ip_address(x)) for x in ips] else: - ip_ints = [int(ipaddress.ip_address(ips))] + ip_ints = [int(ip_address(ips))] await cls.find(cls.latest == True, In(cls.ip_int, ip_ints)).update_many( Set({cls.latest: False}) From bd0ada51e54cc7dc5131391cadab27a21805d882 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 13 Feb 2024 17:24:58 -0500 Subject: [PATCH 035/139] Add new models and fields to SnapshotDoc --- src/cyhy_db/models/snapshot_doc.py | 67 ++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index 5691999..24b3f57 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -1,16 +1,83 @@ # Standard Python Libraries from datetime import datetime +from ipaddress import IPv4Network # Third-Party Libraries from beanie import Document from pydantic import Field from pymongo import ASCENDING, IndexModel +from datetime import datetime +from pydantic import BaseModel, Field +from typing import List, Dict +from bson import ObjectId + + +class VulnerabilityCounts(BaseModel): + critical: int = 0 + high: int = 0 + medium: int = 0 + low: int = 0 + total: int = 0 + + +class WorldData(BaseModel): + host_count: int = 0 + vulnerable_host_count: int = 0 + vulnerabilities: VulnerabilityCounts = Field(default_factory=VulnerabilityCounts) + unique_vulnerabilities: VulnerabilityCounts = Field( + default_factory=VulnerabilityCounts + ) + cvss_average_all: float = 0.0 + cvss_average_vulnerable: float = 0.0 + + +class TicketMetrics(BaseModel): + median: int = 0 + max: int = 0 + + +class TicketOpenMetrics(BaseModel): + # Numbers in this section refer to how long open tix were open AT this date/time + tix_open_as_of_date: datetime = Field(default_factory=datetime.utcnow) + critical: TicketMetrics = Field(default_factory=TicketMetrics) + high: TicketMetrics = Field(default_factory=TicketMetrics) + medium: TicketMetrics = Field(default_factory=TicketMetrics) + low: TicketMetrics = Field(default_factory=TicketMetrics) + + +class TicketCloseMetrics(BaseModel): + # Numbers in this section only include tix that closed AT/AFTER this date/time + tix_closed_after_date: datetime = Field(default_factory=datetime.utcnow) + critical: TicketMetrics = Field(default_factory=TicketMetrics) + high: TicketMetrics = Field(default_factory=TicketMetrics) + medium: TicketMetrics = Field(default_factory=TicketMetrics) + low: TicketMetrics = Field(default_factory=TicketMetrics) class SnapshotDoc(Document): owner: str = Field(...) + descendants_included: List[str] = Field(default=[]) + last_change: datetime = Field(default_factory=datetime.utcnow) start_time: datetime = Field(...) end_time: datetime = Field(...) + latest: bool = Field(default=True) + port_count: int = Field(default=0) + unique_port_count: int = Field(default=0) + unique_operating_systems: int = Field(default=0) + host_count: int = Field(default=0) + vulnerable_host_count: int = Field(default=0) + vulnerabilities: VulnerabilityCounts = Field(default_factory=VulnerabilityCounts) + unique_vulnerabilities: VulnerabilityCounts = Field( + default_factory=VulnerabilityCounts + ) + cvss_average_all: float = Field(default=0.0) + cvss_average_vulnerable: float = Field(default=0.0) + world: WorldData = Field(default_factory=WorldData) + networks: List[IPv4Network] = Field(default=[]) + addresses_scanned: int = Field(default=0) + services: Dict = Field(default_factory=dict) + tix_msec_open: TicketOpenMetrics = Field(default_factory=TicketOpenMetrics) + tix_msec_to_close: TicketCloseMetrics = Field(default_factory=TicketCloseMetrics) class Settings: # Beanie settings From 08507250840d73a686e614ff3e4c1138d7b8b9c7 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:29:53 -0500 Subject: [PATCH 036/139] Reorder model imports to avoid circular references --- src/cyhy_db/models/__init__.py | 34 +++++++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py index 723a1de..00574a4 100644 --- a/src/cyhy_db/models/__init__.py +++ b/src/cyhy_db/models/__init__.py @@ -1,7 +1,35 @@ +# Scan documents (order matters) +from .scan_doc import ScanDoc +from .host_scan_doc import HostScanDoc +from .port_scan_doc import PortScanDoc +from .vuln_scan_doc import VulnScanDoc + +# Snapshot documents (order matters) +from .snapshot_doc import SnapshotDoc +from .report_doc import ReportDoc + +# Other documents from .cve import CVE from .host_doc import HostDoc +from .kev_doc import KEVDoc +from .place_doc import PlaceDoc from .request_doc import RequestDoc -from .scan_doc import ScanDoc -from .snapshot_doc import SnapshotDoc +from .system_control_doc import SystemControlDoc +from .tally_doc import TallyDoc + -__all__ = ["CVE", "HostDoc", "RequestDoc", "ScanDoc", "SnapshotDoc"] +__all__ = [ + "CVE", + "HostDoc", + "HostScanDoc", + "KEVDoc", + "PlaceDoc", + "PortScanDoc", + "RequestDoc", + "ReportDoc", + "ScanDoc", + "SnapshotDoc", + "SystemControlDoc", + "TallyDoc", + "VulnScanDoc", +] From 8feda1d6b5a088314ebbd16dbba6a78e348b489f Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:30:10 -0500 Subject: [PATCH 037/139] Add Protocol enum for TCP and UDP --- src/cyhy_db/models/enum.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index 681a4c9..29f749b 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -41,6 +41,11 @@ class PocType(Enum): TECHNICAL = "TECHNICAL" +class Protocol(Enum): + TCP = "tcp" + UDP = "udp" + + class ReportPeriod(Enum): MONTHLY = "MONTHLY" QUARTERLY = "QUARTERLY" From adc11c7321d135582bd82dfcd64ed02d980f06d8 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:30:36 -0500 Subject: [PATCH 038/139] Add deprecated decorator to utils --- src/cyhy_db/utils/__init__.py | 4 +++- src/cyhy_db/utils/decorators.py | 17 +++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) create mode 100644 src/cyhy_db/utils/decorators.py diff --git a/src/cyhy_db/utils/__init__.py b/src/cyhy_db/utils/__init__.py index a9bfaf4..9cdfc87 100644 --- a/src/cyhy_db/utils/__init__.py +++ b/src/cyhy_db/utils/__init__.py @@ -1,5 +1,7 @@ """Utility functions for cyhy_db.""" +from .decorators import deprecated from .time import utcnow -__all__ = ["utcnow"] + +__all__ = ["deprecated", "utcnow"] diff --git a/src/cyhy_db/utils/decorators.py b/src/cyhy_db/utils/decorators.py new file mode 100644 index 0000000..ef3b8c8 --- /dev/null +++ b/src/cyhy_db/utils/decorators.py @@ -0,0 +1,17 @@ +import warnings + + +def deprecated(reason): + def decorator(func): + if isinstance(reason, str): + message = f"{func.__name__} is deprecated and will be removed in a future version. {reason}" + else: + message = f"{func.__name__} is deprecated and will be removed in a future version." + + def wrapper(*args, **kwargs): + warnings.warn(message, DeprecationWarning, stacklevel=2) + return func(*args, **kwargs) + + return wrapper + + return decorator From 8420216d1c7822a28e730e33c23b3558f73d4d34 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:31:41 -0500 Subject: [PATCH 039/139] Refactor and partially implement helpers with deprecation --- src/cyhy_db/models/host_doc.py | 62 ++++++++++++++++++++++++++++++++-- 1 file changed, 60 insertions(+), 2 deletions(-) diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 4266008..5c4c8b4 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -6,6 +6,7 @@ import random from .enum import Stage, Status from ipaddress import IPv4Address, ip_address +from ..utils import deprecated, utcnow class State(BaseModel): @@ -17,7 +18,7 @@ class HostDoc(Document): id: int = Field(...) # IP address as an integer ip: IPv4Address = Field(...) owner: str = Field(...) - last_change: datetime = Field(default_factory=datetime.utcnow) + last_change: datetime = Field(default_factory=utcnow) next_scan: Optional[datetime] = Field(default=None) state: State = Field(default_factory=lambda: State(reason="new", up=False)) stage: Stage = Field(default=Stage.NETSCAN1) @@ -36,9 +37,10 @@ def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: @before_event(Insert, Replace, ValidateOnSave) async def before_save(self): - self.last_change = datetime.utcnow() + self.last_change = utcnow() class Settings: + # Beanie settings name = "hosts" indexes = [ IndexModel( @@ -88,3 +90,59 @@ class Settings: name="latest_scan_done", ), ] + + def set_state(self, nmap_says_up, has_open_ports, reason=None): + """Sets state.up based on different stage + evidence. nmap has a concept of up which is + different from our definition. An nmap "up" just + means it got a reply, not that there are any open + ports. Note either argument can be None.""" + + if has_open_ports == True: # Only PORTSCAN sends in has_open_ports + self.state = State(True, "open-port") + elif has_open_ports == False: + self.state = State(False, "no-open") + elif nmap_says_up == False: # NETSCAN says host is down + self.state = State(False, reason) + + # TODO: There are a lot of functions in the Python 2 version that may or may not be used. + # Instead of porting them all over, we should just port them as they are needed. + # And rewrite things that can be done better in Python 3. + + # @classmethod + # async def get_count(cls, owner: str, stage: Stage, status: Status): + # return await cls.count( + # cls.owner == owner, cls.stage == stage, cls.status == status + # ) + + @classmethod + @deprecated("Use HostDoc.find_one(HostDoc.ip == ip) instead.") + async def get_by_ip(cls, ip: IPv4Address): + return await cls.find_one(cls.ip == ip) + + # @classmethod + # @deprecated("Use HostDoc.find_one(HostDoc.ip == ip).owner instead.") + # async def get_owner_of_ip(cls, ip: IPv4Address): + # host = await cls.get_by_ip(ip) + # return host.owner + + # @classmethod + # async def get_some_for_stage( + # cls, + # stage: Stage, + # count: int, + # owner: Optional[str] = None, + # waiting: bool = False, + # ): + # if waiting: + # status = {"$in": [Status.READY, Status.WAITING]} + # else: + # status = Status.READY + + # query = cls.find(cls.status == status, cls.stage == stage) + # if owner is not None: + # query = query.find(cls.owner == owner) + + # # Sorting and limiting the results + # results = await query.sort([("priority", 1), ("r", 1)]).limit(count).to_list() + # return results From 8a1cc07d18338f39c99784637fee6d4c4be79a19 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:31:52 -0500 Subject: [PATCH 040/139] Add HostScanDoc model with settings --- src/cyhy_db/models/host_scan_doc.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) create mode 100644 src/cyhy_db/models/host_scan_doc.py diff --git a/src/cyhy_db/models/host_scan_doc.py b/src/cyhy_db/models/host_scan_doc.py new file mode 100644 index 0000000..27b6a53 --- /dev/null +++ b/src/cyhy_db/models/host_scan_doc.py @@ -0,0 +1,20 @@ +from . import ScanDoc +from typing import List +from pymongo import ASCENDING, IndexModel + + +class HostScanDoc(ScanDoc): + name: str + accuracy: int + line: int + classes: List[dict] = [] + + class Settings: + # Beanie settings + name = "host_scans" + indexes = [ + IndexModel( + [("latest", ASCENDING), ("owner", ASCENDING)], name="latest_owner" + ), + IndexModel([("owner", ASCENDING)], name="owner"), + ] From c4d22c6504f74a97b95e9d9566a15a81c3b6aa1b Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:32:08 -0500 Subject: [PATCH 041/139] Add KEVDoc model --- src/cyhy_db/models/kev_doc.py | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 src/cyhy_db/models/kev_doc.py diff --git a/src/cyhy_db/models/kev_doc.py b/src/cyhy_db/models/kev_doc.py new file mode 100644 index 0000000..66548eb --- /dev/null +++ b/src/cyhy_db/models/kev_doc.py @@ -0,0 +1,9 @@ +from beanie import Document + + +class KEVDoc(Document): + id: str # CVE + known_ransomware: bool + + class Settings: + name = "kevs" From bb225978b633914a099ff483d857a999a56a92c9 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:32:21 -0500 Subject: [PATCH 042/139] Add PlaceDoc model --- src/cyhy_db/models/place_doc.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 src/cyhy_db/models/place_doc.py diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py new file mode 100644 index 0000000..1c849b4 --- /dev/null +++ b/src/cyhy_db/models/place_doc.py @@ -0,0 +1,25 @@ +from beanie import Document +from typing import Optional +from pydantic import Field + + +class PlaceDoc(Document): + id: int # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html + name: str + clazz: str = Field(alias="class") # 'class' is a reserved keyword in Python + state: str + state_fips: str + state_name: str + county: Optional[str] = None + county_fips: Optional[str] = None + country: str + country_name: str + latitude_dms: Optional[str] = None + longitude_dms: Optional[str] = None + latitude_dec: float + longitude_dec: float + elevation_meters: Optional[int] = None + elevation_feet: Optional[int] = None + + class Settings: + name = "places" From 0e14f8eb61445e7f555109df8dfb1a2778ff12c1 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:32:36 -0500 Subject: [PATCH 043/139] Add PortScanDoc model with indexes --- src/cyhy_db/models/port_scan_doc.py | 35 +++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 src/cyhy_db/models/port_scan_doc.py diff --git a/src/cyhy_db/models/port_scan_doc.py b/src/cyhy_db/models/port_scan_doc.py new file mode 100644 index 0000000..ac73b90 --- /dev/null +++ b/src/cyhy_db/models/port_scan_doc.py @@ -0,0 +1,35 @@ +from typing import Dict +from pymongo import ASCENDING, IndexModel + +from . import ScanDoc +from .enum import Protocol + + +class PortScanDoc(ScanDoc): + protocol: Protocol + port: int + service: Dict = {} # Assuming no specific structure for "service" + state: str + reason: str + + class Settings: + # Beanie settings + name = "port_scans" + indexes = [ + IndexModel( + [("latest", ASCENDING), ("owner", ASCENDING), ("state", ASCENDING)], + name="latest_owner_state", + ), + IndexModel( + [("latest", ASCENDING), ("service.name", ASCENDING)], + name="latest_service_name", + ), + IndexModel( + [("latest", ASCENDING), ("time", ASCENDING)], + name="latest_time", + ), + IndexModel( + [("owner", ASCENDING)], + name="owner", + ), + ] From d6390d5a830fd25bb6289121391f9929d855b1f0 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:32:48 -0500 Subject: [PATCH 044/139] Add ReportDoc model --- src/cyhy_db/models/report_doc.py | 33 ++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 src/cyhy_db/models/report_doc.py diff --git a/src/cyhy_db/models/report_doc.py b/src/cyhy_db/models/report_doc.py new file mode 100644 index 0000000..736018e --- /dev/null +++ b/src/cyhy_db/models/report_doc.py @@ -0,0 +1,33 @@ +from beanie import Document, Indexed, Link +from datetime import datetime +from typing import List +from pymongo import ASCENDING, IndexModel +from pydantic import Field + +from . import SnapshotDoc +from .enum import ReportType +from ..utils import utcnow + + +class ReportDoc(Document): + owner: str + generated_time: datetime = Field(default_factory=utcnow) + snapshots: List[Link[SnapshotDoc]] + report_types: List[ReportType] + + class Settings: + name = "reports" + indexes = [ + IndexModel( + [ + ("owner", ASCENDING), + ], + name="owner", + ), + IndexModel( + [ + ("generated_time", ASCENDING), + ], + name="generated_time", + ), + ] From 3102b9cafff6252d8d7e377a95db80d5bfdd6dbc Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:33:47 -0500 Subject: [PATCH 045/139] Replace deprecated datetime utcnow calls --- src/cyhy_db/models/request_doc.py | 6 ++++-- src/cyhy_db/models/snapshot_doc.py | 8 +++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 277085b..8bbdf7b 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -19,6 +19,8 @@ Stage, ) +from ..utils import utcnow + BOGUS_ID = "bogus_id_replace_me" @@ -71,11 +73,11 @@ class RequestDoc(Document): id: str = Field(default=BOGUS_ID) agency: Agency children: List[Link["RequestDoc"]] = Field(default=[]) - enrolled: datetime = Field(default_factory=datetime.utcnow) + enrolled: datetime = Field(default_factory=utcnow) init_stage: Stage = Field(default=Stage.NETSCAN1) key: Optional[str] = Field(default=None) networks: List[IPv4Network] = Field(default=[]) - period_start: datetime = Field(default_factory=datetime.utcnow) + period_start: datetime = Field(default_factory=utcnow) report_period: ReportPeriod = Field(default=ReportPeriod.WEEKLY) report_types: List[ReportType] = Field(default=[]) retired: bool = False diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index 24b3f57..daeeadd 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -11,6 +11,8 @@ from typing import List, Dict from bson import ObjectId +from ..utils import utcnow + class VulnerabilityCounts(BaseModel): critical: int = 0 @@ -38,7 +40,7 @@ class TicketMetrics(BaseModel): class TicketOpenMetrics(BaseModel): # Numbers in this section refer to how long open tix were open AT this date/time - tix_open_as_of_date: datetime = Field(default_factory=datetime.utcnow) + tix_open_as_of_date: datetime = Field(default_factory=utcnow) critical: TicketMetrics = Field(default_factory=TicketMetrics) high: TicketMetrics = Field(default_factory=TicketMetrics) medium: TicketMetrics = Field(default_factory=TicketMetrics) @@ -47,7 +49,7 @@ class TicketOpenMetrics(BaseModel): class TicketCloseMetrics(BaseModel): # Numbers in this section only include tix that closed AT/AFTER this date/time - tix_closed_after_date: datetime = Field(default_factory=datetime.utcnow) + tix_closed_after_date: datetime = Field(default_factory=utcnow) critical: TicketMetrics = Field(default_factory=TicketMetrics) high: TicketMetrics = Field(default_factory=TicketMetrics) medium: TicketMetrics = Field(default_factory=TicketMetrics) @@ -57,7 +59,7 @@ class TicketCloseMetrics(BaseModel): class SnapshotDoc(Document): owner: str = Field(...) descendants_included: List[str] = Field(default=[]) - last_change: datetime = Field(default_factory=datetime.utcnow) + last_change: datetime = Field(default_factory=utcnow) start_time: datetime = Field(...) end_time: datetime = Field(...) latest: bool = Field(default=True) From 390a15377907378f603ea8b0d81be3969900dc74 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:33:58 -0500 Subject: [PATCH 046/139] Add SystemControlDoc model --- src/cyhy_db/models/system_control_doc.py | 35 ++++++++++++++++++++++++ 1 file changed, 35 insertions(+) create mode 100644 src/cyhy_db/models/system_control_doc.py diff --git a/src/cyhy_db/models/system_control_doc.py b/src/cyhy_db/models/system_control_doc.py new file mode 100644 index 0000000..61f8df1 --- /dev/null +++ b/src/cyhy_db/models/system_control_doc.py @@ -0,0 +1,35 @@ +from beanie import Document +from datetime import datetime +from pydantic import Field +from typing import Optional +import asyncio + +from .enum import ControlAction, ControlTarget +from ..utils import utcnow + +CONTROL_DOC_POLL_INTERVAL = 5 # seconds + + +class SystemControlDoc(Document): + action: ControlAction + sender: str # Free-form, for UI / Logging + target: ControlTarget + reason: str # Free-form, for UI / Logging + time: datetime = Field(default_factory=utcnow) # creation time + completed: bool = False # Set to True when after the action has occurred + + class Settings: + name = "control" + + @classmethod + async def wait_for_completion(cls, document_id, timeout: Optional[int] = None): + """Wait for this control action to complete. If a timeout is set, only wait a maximum of timeout seconds. + Returns True if the document was completed, False otherwise.""" + start_time = utcnow() + while True: + doc = await cls.get(document_id) + if doc.completed: + return True + if timeout and (utcnow() - start_time).total_seconds() > timeout: + return False + await asyncio.sleep(CONTROL_DOC_POLL_INTERVAL) From 230ec5990850f20dfea306d8ee7978ba7875d5c5 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:34:10 -0500 Subject: [PATCH 047/139] Add TallyDoc model --- src/cyhy_db/models/tally_doc.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) create mode 100644 src/cyhy_db/models/tally_doc.py diff --git a/src/cyhy_db/models/tally_doc.py b/src/cyhy_db/models/tally_doc.py new file mode 100644 index 0000000..6af6a2c --- /dev/null +++ b/src/cyhy_db/models/tally_doc.py @@ -0,0 +1,33 @@ +from beanie import Document, before_event, Insert, Replace, ValidateOnSave +from datetime import datetime +from pydantic import BaseModel, Field + +from ..utils import utcnow + + +class StatusCounts(BaseModel): + READY: int = 0 + WAITING: int = 0 + DONE: int = 0 + RUNNING: int = 0 + + +class Counts(BaseModel): + PORTSCAN: StatusCounts = Field(default_factory=StatusCounts) + BASESCAN: StatusCounts = Field(default_factory=StatusCounts) + VULNSCAN: StatusCounts = Field(default_factory=StatusCounts) + NETSCAN1: StatusCounts = Field(default_factory=StatusCounts) + NETSCAN2: StatusCounts = Field(default_factory=StatusCounts) + + +class TallyDoc(Document): + _id: str # owner_id + counts: Counts = Field(default_factory=Counts) + last_change: datetime = Field(default_factory=utcnow) + + @before_event(Insert, Replace, ValidateOnSave) + async def before_save(self): + self.last_change = utcnow() + + class Settings: + name = "tallies" From 7bb64650e6964b48828baf210dcc9689315bfda0 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:34:23 -0500 Subject: [PATCH 048/139] Add VulnScanDoc model --- src/cyhy_db/models/vuln_scan_doc.py | 36 +++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 src/cyhy_db/models/vuln_scan_doc.py diff --git a/src/cyhy_db/models/vuln_scan_doc.py b/src/cyhy_db/models/vuln_scan_doc.py new file mode 100644 index 0000000..3973e98 --- /dev/null +++ b/src/cyhy_db/models/vuln_scan_doc.py @@ -0,0 +1,36 @@ +from typing import Dict +from pymongo import ASCENDING, IndexModel +from datetime import datetime + +from . import ScanDoc +from .enum import Protocol + + +class VulnScanDoc(ScanDoc): + protocol: Protocol + port: int + service: str + cvss_base_score: float + cvss_vector: str + description: str + fname: str + plugin_family: str + plugin_id: int + plugin_modification_date: datetime + plugin_name: str + plugin_publication_date: datetime + plugin_type: str + risk_factor: str + severity: int + solution: str + synopsis: str + + class Settings: + # Beanie settings + name = "vuln_scans" + indexes = [ + IndexModel( + [("owner", ASCENDING), ("latest", ASCENDING), ("severity", ASCENDING)], + name="owner_latest_severity", + ), + ] From 345fb139c1ef92c5f3f05a62eba2de8f19bad5c8 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:34:37 -0500 Subject: [PATCH 049/139] Update import statements in db.py --- src/cyhy_db/db.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index 9698924..9357d41 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -2,12 +2,24 @@ from beanie import init_beanie from motor.motor_asyncio import AsyncIOMotorClient -from .models.cve import CVE -from .models.request_doc import RequestDoc -from .models.scan_doc import ScanDoc -from .models.snapshot_doc import SnapshotDoc +from .models import * -ALL_MODELS = [CVE, RequestDoc, ScanDoc, SnapshotDoc] + +ALL_MODELS = [ + CVE, + HostDoc, + HostScanDoc, + KEVDoc, + PlaceDoc, + PortScanDoc, + RequestDoc, + ReportDoc, + ScanDoc, + SnapshotDoc, + SystemControlDoc, + TallyDoc, + VulnScanDoc, +] async def initialize_db(db_uri: str, db_name: str) -> None: From c2d40eb3c5297911ceae50db259e136816210f66 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:35:09 -0500 Subject: [PATCH 050/139] Replace deprecated utcnow calls --- tests/test_scan_doc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index 1ca4b46..a24676f 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -10,6 +10,7 @@ # cisagov Libraries from cyhy_db.models import ScanDoc, SnapshotDoc +from cyhy_db.utils import utcnow VALID_IP_1_STR = "0.0.0.1" VALID_IP_2_STR = "0.0.0.2" @@ -134,8 +135,8 @@ async def test_tag_latest(): owner = "TAG_LATEST" snapshot_doc = SnapshotDoc( owner=owner, - start_time=datetime.datetime.utcnow(), - end_time=datetime.datetime.utcnow(), + start_time=utcnow(), + end_time=utcnow(), ) await snapshot_doc.save() # Create a ScanDoc object From e71cd08463a6de0f37a1a5872dc8cfbeb611cd36 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:35:33 -0500 Subject: [PATCH 051/139] Refactor test_find to test_get_by_ip --- tests/test_host_doc.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/test_host_doc.py b/tests/test_host_doc.py index a1e0a57..413504a 100644 --- a/tests/test_host_doc.py +++ b/tests/test_host_doc.py @@ -26,8 +26,7 @@ async def test_save(): assert host_doc.id == VALID_IP_1_INT -async def test_find(): +async def test_get_by_ip(): # Find a HostDoc object by its IP address - host_doc = await HostDoc.find_one(HostDoc.id == VALID_IP_1_INT) - print(host_doc) + host_doc = await HostDoc.get_by_ip(ip_address(VALID_IP_1_STR)) assert host_doc.ip == ip_address(VALID_IP_1_STR) From a59a9219cd1f2c7b3d5c2267413299ff9113e2cb Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 16 Feb 2024 14:35:44 -0500 Subject: [PATCH 052/139] Add test data generator for CVE and RequestDoc models --- tests/test_data_generator.py | 158 +++++++++++++++++++++++++++++++++++ 1 file changed, 158 insertions(+) create mode 100644 tests/test_data_generator.py diff --git a/tests/test_data_generator.py b/tests/test_data_generator.py new file mode 100644 index 0000000..0a39d8d --- /dev/null +++ b/tests/test_data_generator.py @@ -0,0 +1,158 @@ +from cyhy_db.models import CVE, RequestDoc +from cyhy_db.models.request_doc import Agency, Contact, Location, Window +from cyhy_db.models.enum import ( + AgencyType, + DayOfWeek, + CVSSVersion, + PocType, + ReportPeriod, + ReportType, + ScanType, + Scheduler, + Stage, +) +from cyhy_db.utils import utcnow +from mimesis_factory import MimesisField +from mimesis.locales import DEFAULT_LOCALE +from mimesis.providers.base import BaseProvider +from pytest_factoryboy import register +import factory +from mimesis import Generic +import ipaddress +import random +from datetime import datetime + + +class CyHyProvider(BaseProvider): + class Meta: + name = "cyhy_provider" + + def cve_id(self, year=None): + # If year is None, generate a random year between 1999 and the current year + if year is None: + year = self.random.randint(1999, datetime.now().year) + + # Generate a random number for the CVE, ensuring it has a leading zero if necessary + number = self.random.randint(1, 99999) + + return f"CVE-{year}-{number:05d}" + + def network_ipv4(self): + # Generate a base IP address + base_ip = generic.internet.ip_v4() + # Choose a random CIDR between 24-30 to ensure a smaller network size and avoid host bits set error + cidr = random.randint(24, 30) + # Create the network address + network = ipaddress.IPv4Network(f"{base_ip}/{cidr}", strict=False) + return network + + +generic = Generic(locale=DEFAULT_LOCALE) +generic.add_provider(CyHyProvider) + + +@register +class CVEFactory(factory.Factory): + class Meta(object): + model = CVE + + id = factory.LazyFunction(lambda: generic.cyhy_provider.cve_id()) + cvss_score = factory.LazyFunction(lambda: round(random.uniform(0, 10), 1)) + cvss_version = factory.LazyFunction(lambda: random.choice(list(CVSSVersion))) + severity = factory.LazyFunction(lambda: random.randint(1, 4)) + + +class AgencyFactory(factory.Factory): + class Meta(object): + model = Agency + + name = factory.Faker("company") + # Generate an acronym from the name + acronym = factory.LazyAttribute( + lambda o: "".join(word[0].upper() for word in o.name.split()) + ) + type = factory.LazyFunction(lambda: random.choice(list(AgencyType))) + contacts = factory.LazyFunction( + lambda: [ContactFactory() for _ in range(random.randint(1, 5))] + ) + location = factory.LazyFunction(lambda: LocationFactory()) + + +class ContactFactory(factory.Factory): + class Meta(object): + model = Contact + + email = factory.Faker("email") + name = factory.Faker("name") + phone = factory.Faker("phone_number") + type = factory.LazyFunction(lambda: random.choice(list(PocType))) + + +class LocationFactory(factory.Factory): + class Meta(object): + model = Location + + country_name = factory.Faker("country") + country = factory.Faker("country_code") + county_fips = factory.Faker("numerify", text="##") + county = factory.Faker("city") + gnis_id = factory.Faker("numerify", text="#######") + name = factory.Faker("city") + state_fips = factory.Faker("numerify", text="##") + state_name = factory.Faker("state") + state = factory.Faker("state_abbr") + + +class WindowFactory(factory.Factory): + class Meta(object): + model = Window + + day = factory.LazyFunction(lambda: random.choice(list(DayOfWeek))) + duration = factory.LazyFunction(lambda: random.randint(0, 168)) + start = factory.Faker("time", pattern="%H:%M:%S") + + +class RequestDocFactory(factory.Factory): + class Meta(object): + model = RequestDoc + + id = factory.LazyAttribute( + lambda o: o.agency.acronym + "-" + str(random.randint(1, 1000)) + ) + agency = factory.SubFactory(AgencyFactory) + enrolled = factory.LazyFunction(utcnow) + init_stage = factory.LazyFunction(lambda: random.choice(list(Stage))) + key = factory.Faker("password") + period_start = factory.LazyFunction(utcnow) + report_period = factory.LazyFunction(lambda: random.choice(list(ReportPeriod))) + retired = factory.LazyFunction(lambda: random.choice([True, False])) + scheduler = factory.LazyFunction(lambda: random.choice(list(Scheduler))) + stakeholder = factory.LazyFunction(lambda: random.choice([True, False])) + windows = factory.LazyFunction( + lambda: [WindowFactory() for _ in range(random.randint(1, 5))] + ) + networks = factory.LazyFunction( + lambda: [ + generic.cyhy_provider.network_ipv4() for _ in range(random.randint(1, 5)) + ] + ) + # create a set of 1 to 3 random scan types from the ScanType enum + scan_types = factory.LazyFunction( + lambda: set( + [random.choice(list(ScanType)) for _ in range(random.randint(1, 3))] + ) + ) + + +async def test_create_cves(): + for _ in range(100): + cve = CVEFactory() + print(cve) + await cve.save() + + +async def test_create_request_docs(): + for _ in range(100): + request_doc = RequestDocFactory() + print(request_doc) + await request_doc.save() From 12e7047a6b79b2524413124ac0dec231ec1aab51 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 10:42:07 -0400 Subject: [PATCH 053/139] Add .hypothesis and .vscode to .gitignore --- .gitignore | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.gitignore b/.gitignore index 242b4aa..82ac921 100644 --- a/.gitignore +++ b/.gitignore @@ -5,8 +5,12 @@ ## Python ## __pycache__ .coverage +.hypothesis .mypy_cache .pytest_cache .python-version *.egg-info dist + +## VSCode ## +.vscode From c43cf3df9828397d38f225b06bfc9dea8a5687de Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 13:05:13 -0400 Subject: [PATCH 054/139] Add note about pydantic extras --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 37f74f2..589ac3f 100644 --- a/setup.py +++ b/setup.py @@ -93,7 +93,7 @@ def get_version(version_file): install_requires=[ "beanie == 1.25.0", "docopt == 0.6.2", - "pydantic[email, hypothesis]", + "pydantic[email, hypothesis]", # hypothesis plugin is currently disabled: https://github.com/pydantic/pydantic/issues/4682 "schema == 0.7.5", "setuptools >= 69.0.3", ], From b2574f6c00476049cb6d276cbdc76dd7f806eb22 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 13:06:34 -0400 Subject: [PATCH 055/139] Add NotificationDoc model to cyhy_db --- src/cyhy_db/db.py | 1 + src/cyhy_db/models/__init__.py | 2 ++ src/cyhy_db/models/notification_doc.py | 17 +++++++++++++++++ 3 files changed, 20 insertions(+) create mode 100644 src/cyhy_db/models/notification_doc.py diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index 9357d41..faa34e7 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -10,6 +10,7 @@ HostDoc, HostScanDoc, KEVDoc, + NotificationDoc, PlaceDoc, PortScanDoc, RequestDoc, diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py index 00574a4..05748be 100644 --- a/src/cyhy_db/models/__init__.py +++ b/src/cyhy_db/models/__init__.py @@ -12,6 +12,7 @@ from .cve import CVE from .host_doc import HostDoc from .kev_doc import KEVDoc +from .notification_doc import NotificationDoc from .place_doc import PlaceDoc from .request_doc import RequestDoc from .system_control_doc import SystemControlDoc @@ -23,6 +24,7 @@ "HostDoc", "HostScanDoc", "KEVDoc", + "NotificationDoc", "PlaceDoc", "PortScanDoc", "RequestDoc", diff --git a/src/cyhy_db/models/notification_doc.py b/src/cyhy_db/models/notification_doc.py new file mode 100644 index 0000000..82dd405 --- /dev/null +++ b/src/cyhy_db/models/notification_doc.py @@ -0,0 +1,17 @@ +from beanie import Document, BeanieObjectId +from pydantic import Field, BaseModel, ConfigDict +from bson import ObjectId +from typing import List + + +class NotificationDoc(Document): + model_config = ConfigDict(extra="forbid") + + ticket_id: BeanieObjectId = Field(...) # ticket id that triggered the notification + ticket_owner: str # owner of the ticket + generated_for: List[str] = Field( + default=[] + ) # list of owners built as notifications are generated + + class Settings: + name = "notifications" From b427576f538438c08c1e7042295a16af50bb8a35 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 13:08:19 -0400 Subject: [PATCH 056/139] Update model configurations to use ConfigDict instead of Config class --- src/cyhy_db/models/cve.py | 10 ++++------ src/cyhy_db/models/host_doc.py | 9 +++++---- src/cyhy_db/models/host_scan_doc.py | 3 +++ src/cyhy_db/models/kev_doc.py | 3 +++ src/cyhy_db/models/place_doc.py | 4 +++- src/cyhy_db/models/port_scan_doc.py | 2 ++ src/cyhy_db/models/report_doc.py | 4 +++- src/cyhy_db/models/request_doc.py | 14 +++++++++++++- src/cyhy_db/models/scan_doc.py | 10 ++++------ src/cyhy_db/models/snapshot_doc.py | 14 +++++++++++++- src/cyhy_db/models/system_control_doc.py | 4 +++- src/cyhy_db/models/tally_doc.py | 16 +++++++++++----- src/cyhy_db/models/vuln_scan_doc.py | 3 +++ 13 files changed, 70 insertions(+), 26 deletions(-) diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 0c0489c..7878163 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -3,12 +3,15 @@ # Third-Party Libraries from beanie import Document, Indexed, ValidateOnSave, before_event -from pydantic import Field, model_validator +from pydantic import Field, model_validator, ConfigDict from .enum import CVSSVersion class CVE(Document): + # Validate on assignment so ip_int is recalculated as ip is set + model_config = ConfigDict(extra="forbid", validate_assignment=True) + id: str = Indexed(primary_field=True) # CVE ID cvss_score: float = Field(ge=0.0, le=10.0) cvss_version: CVSSVersion = Field(default=CVSSVersion.V3_1) @@ -36,11 +39,6 @@ def calculate_severity(cls, values: Dict[str, Any]) -> Dict[str, Any]: values["severity"] = 1 return values - class Config: - # Pydantic configuration - # Validate on assignment so ip_int is recalculated as ip is set - validate_assignment = True - class Settings: # Beanie settings name = "cves" diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 5c4c8b4..4c665c5 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -1,6 +1,6 @@ from beanie import Document, before_event, Indexed, Insert, Replace, ValidateOnSave from datetime import datetime -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field, model_validator, ConfigDict from pymongo import ASCENDING, IndexModel from typing import Any, Dict, Optional, Tuple import random @@ -15,7 +15,9 @@ class State(BaseModel): class HostDoc(Document): - id: int = Field(...) # IP address as an integer + model_config = ConfigDict(extra="forbid") + + id: int = Field() # IP address as an integer ip: IPv4Address = Field(...) owner: str = Field(...) last_change: datetime = Field(default_factory=utcnow) @@ -31,8 +33,7 @@ class HostDoc(Document): @model_validator(mode="before") def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: # ip may still be string if it was just set - values["id"] = int(ip_address(values["ip"])) - print(values) + values["_id"] = int(ip_address(values["ip"])) return values @before_event(Insert, Replace, ValidateOnSave) diff --git a/src/cyhy_db/models/host_scan_doc.py b/src/cyhy_db/models/host_scan_doc.py index 27b6a53..fc41120 100644 --- a/src/cyhy_db/models/host_scan_doc.py +++ b/src/cyhy_db/models/host_scan_doc.py @@ -1,9 +1,12 @@ from . import ScanDoc from typing import List from pymongo import ASCENDING, IndexModel +from pydantic import ConfigDict class HostScanDoc(ScanDoc): + model_config = ConfigDict(extra="forbid") + name: str accuracy: int line: int diff --git a/src/cyhy_db/models/kev_doc.py b/src/cyhy_db/models/kev_doc.py index 66548eb..0faf085 100644 --- a/src/cyhy_db/models/kev_doc.py +++ b/src/cyhy_db/models/kev_doc.py @@ -1,7 +1,10 @@ from beanie import Document +from pydantic import ConfigDict class KEVDoc(Document): + model_config = ConfigDict(extra="forbid") + id: str # CVE known_ransomware: bool diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index 1c849b4..64fc7da 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -1,9 +1,11 @@ from beanie import Document from typing import Optional -from pydantic import Field +from pydantic import Field, ConfigDict class PlaceDoc(Document): + model_config = ConfigDict(extra="forbid") + id: int # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html name: str clazz: str = Field(alias="class") # 'class' is a reserved keyword in Python diff --git a/src/cyhy_db/models/port_scan_doc.py b/src/cyhy_db/models/port_scan_doc.py index ac73b90..62b307d 100644 --- a/src/cyhy_db/models/port_scan_doc.py +++ b/src/cyhy_db/models/port_scan_doc.py @@ -1,11 +1,13 @@ from typing import Dict from pymongo import ASCENDING, IndexModel +from pydantic import ConfigDict from . import ScanDoc from .enum import Protocol class PortScanDoc(ScanDoc): + model_config = ConfigDict(extra="forbid") protocol: Protocol port: int service: Dict = {} # Assuming no specific structure for "service" diff --git a/src/cyhy_db/models/report_doc.py b/src/cyhy_db/models/report_doc.py index 736018e..ebe5605 100644 --- a/src/cyhy_db/models/report_doc.py +++ b/src/cyhy_db/models/report_doc.py @@ -2,7 +2,7 @@ from datetime import datetime from typing import List from pymongo import ASCENDING, IndexModel -from pydantic import Field +from pydantic import Field, ConfigDict from . import SnapshotDoc from .enum import ReportType @@ -10,6 +10,8 @@ class ReportDoc(Document): + model_config = ConfigDict(extra="forbid") + owner: str generated_time: datetime = Field(default_factory=utcnow) snapshots: List[Link[SnapshotDoc]] diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 8bbdf7b..3437601 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -6,7 +6,7 @@ # Third-Party Libraries from beanie import Document, Insert, Link, Replace, ValidateOnSave, before_event -from pydantic import BaseModel, EmailStr, Field, field_validator +from pydantic import BaseModel, EmailStr, Field, field_validator, ConfigDict from .enum import ( AgencyType, @@ -25,6 +25,8 @@ class Contact(BaseModel): + model_config = ConfigDict(extra="forbid") + email: EmailStr name: str phone: str @@ -32,6 +34,8 @@ class Contact(BaseModel): class Location(BaseModel): + model_config = ConfigDict(extra="forbid") + country_name: str country: str county_fips: str @@ -44,6 +48,8 @@ class Location(BaseModel): class Agency(BaseModel): + model_config = ConfigDict(extra="forbid") + name: str acronym: str type: Optional[AgencyType] = Field(default=None) @@ -52,11 +58,15 @@ class Agency(BaseModel): class ScanLimit(BaseModel): + model_config = ConfigDict(extra="forbid") + scan_type: ScanType = Field(..., alias="scanType") concurrent: int = Field(ge=0) class Window(BaseModel): + model_config = ConfigDict(extra="forbid") + day: DayOfWeek = Field(default=DayOfWeek.SUNDAY) duration: int = Field(default=168, ge=0, le=168) start: str = Field(default="00:00:00") @@ -70,6 +80,8 @@ def validate_start(cls, v): class RequestDoc(Document): + model_config = ConfigDict(extra="forbid") + id: str = Field(default=BOGUS_ID) agency: Agency children: List[Link["RequestDoc"]] = Field(default=[]) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 8421f8f..b1f239a 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -8,11 +8,14 @@ from beanie.operators import In, Push, Set from bson import ObjectId from bson.dbref import DBRef -from pydantic import Field, model_validator +from pydantic import Field, model_validator, ConfigDict from pymongo import ASCENDING, IndexModel class ScanDoc(Document): + # Validate on assignment so ip_int is recalculated as ip is set + model_config = ConfigDict(extra="forbid", validate_assignment=True) + ip: IPv4Address = Field(...) ip_int: int = Field(...) latest: bool = Field(default=True) @@ -27,11 +30,6 @@ def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: values["ip_int"] = int(ip_address(values["ip"])) return values - class Config: - # Pydantic configuration - # Validate on assignment so ip_int is recalculated as ip is set - validate_assignment = True - class Settings: # Beanie settings name = "scandocs" diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index daeeadd..be36fca 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -4,7 +4,7 @@ # Third-Party Libraries from beanie import Document -from pydantic import Field +from pydantic import Field, ConfigDict from pymongo import ASCENDING, IndexModel from datetime import datetime from pydantic import BaseModel, Field @@ -15,6 +15,8 @@ class VulnerabilityCounts(BaseModel): + model_config = ConfigDict(extra="forbid") + critical: int = 0 high: int = 0 medium: int = 0 @@ -23,6 +25,8 @@ class VulnerabilityCounts(BaseModel): class WorldData(BaseModel): + model_config = ConfigDict(extra="forbid") + host_count: int = 0 vulnerable_host_count: int = 0 vulnerabilities: VulnerabilityCounts = Field(default_factory=VulnerabilityCounts) @@ -34,11 +38,15 @@ class WorldData(BaseModel): class TicketMetrics(BaseModel): + model_config = ConfigDict(extra="forbid") + median: int = 0 max: int = 0 class TicketOpenMetrics(BaseModel): + model_config = ConfigDict(extra="forbid") + # Numbers in this section refer to how long open tix were open AT this date/time tix_open_as_of_date: datetime = Field(default_factory=utcnow) critical: TicketMetrics = Field(default_factory=TicketMetrics) @@ -48,6 +56,8 @@ class TicketOpenMetrics(BaseModel): class TicketCloseMetrics(BaseModel): + model_config = ConfigDict(extra="forbid") + # Numbers in this section only include tix that closed AT/AFTER this date/time tix_closed_after_date: datetime = Field(default_factory=utcnow) critical: TicketMetrics = Field(default_factory=TicketMetrics) @@ -57,6 +67,8 @@ class TicketCloseMetrics(BaseModel): class SnapshotDoc(Document): + model_config = ConfigDict(extra="forbid") + owner: str = Field(...) descendants_included: List[str] = Field(default=[]) last_change: datetime = Field(default_factory=utcnow) diff --git a/src/cyhy_db/models/system_control_doc.py b/src/cyhy_db/models/system_control_doc.py index 61f8df1..476125d 100644 --- a/src/cyhy_db/models/system_control_doc.py +++ b/src/cyhy_db/models/system_control_doc.py @@ -1,6 +1,6 @@ from beanie import Document from datetime import datetime -from pydantic import Field +from pydantic import Field, ConfigDict from typing import Optional import asyncio @@ -11,6 +11,8 @@ class SystemControlDoc(Document): + model_config = ConfigDict(extra="forbid") + action: ControlAction sender: str # Free-form, for UI / Logging target: ControlTarget diff --git a/src/cyhy_db/models/tally_doc.py b/src/cyhy_db/models/tally_doc.py index 6af6a2c..5be3ce8 100644 --- a/src/cyhy_db/models/tally_doc.py +++ b/src/cyhy_db/models/tally_doc.py @@ -1,26 +1,32 @@ from beanie import Document, before_event, Insert, Replace, ValidateOnSave from datetime import datetime -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, ConfigDict from ..utils import utcnow class StatusCounts(BaseModel): - READY: int = 0 - WAITING: int = 0 + model_config = ConfigDict(extra="forbid") + DONE: int = 0 + READY: int = 0 RUNNING: int = 0 + WAITING: int = 0 class Counts(BaseModel): - PORTSCAN: StatusCounts = Field(default_factory=StatusCounts) + model_config = ConfigDict(extra="forbid") + BASESCAN: StatusCounts = Field(default_factory=StatusCounts) - VULNSCAN: StatusCounts = Field(default_factory=StatusCounts) NETSCAN1: StatusCounts = Field(default_factory=StatusCounts) NETSCAN2: StatusCounts = Field(default_factory=StatusCounts) + PORTSCAN: StatusCounts = Field(default_factory=StatusCounts) + VULNSCAN: StatusCounts = Field(default_factory=StatusCounts) class TallyDoc(Document): + model_config = ConfigDict(extra="forbid") + _id: str # owner_id counts: Counts = Field(default_factory=Counts) last_change: datetime = Field(default_factory=utcnow) diff --git a/src/cyhy_db/models/vuln_scan_doc.py b/src/cyhy_db/models/vuln_scan_doc.py index 3973e98..65f31ff 100644 --- a/src/cyhy_db/models/vuln_scan_doc.py +++ b/src/cyhy_db/models/vuln_scan_doc.py @@ -1,4 +1,5 @@ from typing import Dict +from pydantic import ConfigDict from pymongo import ASCENDING, IndexModel from datetime import datetime @@ -7,6 +8,8 @@ class VulnScanDoc(ScanDoc): + model_config = ConfigDict(extra="forbid") + protocol: Protocol port: int service: str From 357bc941668ddc31750243c22bbdf64c2297c582 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 13:23:15 -0400 Subject: [PATCH 057/139] Organize imports per isort --- src/cyhy_db/db.py | 1 - src/cyhy_db/models/cve.py | 2 +- src/cyhy_db/models/host_doc.py | 16 +++++---- src/cyhy_db/models/host_scan_doc.py | 8 +++-- src/cyhy_db/models/kev_doc.py | 1 + src/cyhy_db/models/notification_doc.py | 9 +++-- src/cyhy_db/models/place_doc.py | 7 ++-- src/cyhy_db/models/port_scan_doc.py | 5 ++- src/cyhy_db/models/report_doc.py | 9 +++-- src/cyhy_db/models/request_doc.py | 5 ++- src/cyhy_db/models/scan_doc.py | 2 +- src/cyhy_db/models/snapshot_doc.py | 8 ++--- src/cyhy_db/models/system_control_doc.py | 11 +++--- src/cyhy_db/models/tally_doc.py | 7 ++-- src/cyhy_db/models/vuln_scan_doc.py | 5 ++- src/cyhy_db/utils/__init__.py | 1 - src/cyhy_db/utils/decorators.py | 1 + tests/test_data_generator.py | 45 +++++++++++++----------- tests/test_host_doc.py | 5 ++- 19 files changed, 91 insertions(+), 57 deletions(-) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index faa34e7..cfb797a 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -4,7 +4,6 @@ from .models import * - ALL_MODELS = [ CVE, HostDoc, diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 7878163..54ac75e 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -3,7 +3,7 @@ # Third-Party Libraries from beanie import Document, Indexed, ValidateOnSave, before_event -from pydantic import Field, model_validator, ConfigDict +from pydantic import ConfigDict, Field, model_validator from .enum import CVSSVersion diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 4c665c5..3339b70 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -1,12 +1,16 @@ -from beanie import Document, before_event, Indexed, Insert, Replace, ValidateOnSave +# Standard Python Libraries from datetime import datetime -from pydantic import BaseModel, Field, model_validator, ConfigDict -from pymongo import ASCENDING, IndexModel -from typing import Any, Dict, Optional, Tuple -import random -from .enum import Stage, Status from ipaddress import IPv4Address, ip_address +import random +from typing import Any, Dict, Optional, Tuple + +# Third-Party Libraries +from beanie import Document, Indexed, Insert, Replace, ValidateOnSave, before_event +from pydantic import BaseModel, ConfigDict, Field, model_validator +from pymongo import ASCENDING, IndexModel + from ..utils import deprecated, utcnow +from .enum import Stage, Status class State(BaseModel): diff --git a/src/cyhy_db/models/host_scan_doc.py b/src/cyhy_db/models/host_scan_doc.py index fc41120..010f350 100644 --- a/src/cyhy_db/models/host_scan_doc.py +++ b/src/cyhy_db/models/host_scan_doc.py @@ -1,7 +1,11 @@ -from . import ScanDoc +# Standard Python Libraries from typing import List -from pymongo import ASCENDING, IndexModel + +# Third-Party Libraries from pydantic import ConfigDict +from pymongo import ASCENDING, IndexModel + +from . import ScanDoc class HostScanDoc(ScanDoc): diff --git a/src/cyhy_db/models/kev_doc.py b/src/cyhy_db/models/kev_doc.py index 0faf085..20859bf 100644 --- a/src/cyhy_db/models/kev_doc.py +++ b/src/cyhy_db/models/kev_doc.py @@ -1,3 +1,4 @@ +# Third-Party Libraries from beanie import Document from pydantic import ConfigDict diff --git a/src/cyhy_db/models/notification_doc.py b/src/cyhy_db/models/notification_doc.py index 82dd405..922e71f 100644 --- a/src/cyhy_db/models/notification_doc.py +++ b/src/cyhy_db/models/notification_doc.py @@ -1,8 +1,11 @@ -from beanie import Document, BeanieObjectId -from pydantic import Field, BaseModel, ConfigDict -from bson import ObjectId +# Standard Python Libraries from typing import List +# Third-Party Libraries +from beanie import BeanieObjectId, Document +from bson import ObjectId +from pydantic import BaseModel, ConfigDict, Field + class NotificationDoc(Document): model_config = ConfigDict(extra="forbid") diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index 64fc7da..a0ad557 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -1,6 +1,9 @@ -from beanie import Document +# Standard Python Libraries from typing import Optional -from pydantic import Field, ConfigDict + +# Third-Party Libraries +from beanie import Document +from pydantic import ConfigDict, Field class PlaceDoc(Document): diff --git a/src/cyhy_db/models/port_scan_doc.py b/src/cyhy_db/models/port_scan_doc.py index 62b307d..9af2801 100644 --- a/src/cyhy_db/models/port_scan_doc.py +++ b/src/cyhy_db/models/port_scan_doc.py @@ -1,6 +1,9 @@ +# Standard Python Libraries from typing import Dict -from pymongo import ASCENDING, IndexModel + +# Third-Party Libraries from pydantic import ConfigDict +from pymongo import ASCENDING, IndexModel from . import ScanDoc from .enum import Protocol diff --git a/src/cyhy_db/models/report_doc.py b/src/cyhy_db/models/report_doc.py index ebe5605..b917f80 100644 --- a/src/cyhy_db/models/report_doc.py +++ b/src/cyhy_db/models/report_doc.py @@ -1,12 +1,15 @@ -from beanie import Document, Indexed, Link +# Standard Python Libraries from datetime import datetime from typing import List + +# Third-Party Libraries +from beanie import Document, Indexed, Link +from pydantic import ConfigDict, Field from pymongo import ASCENDING, IndexModel -from pydantic import Field, ConfigDict from . import SnapshotDoc -from .enum import ReportType from ..utils import utcnow +from .enum import ReportType class ReportDoc(Document): diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 3437601..9add79e 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -6,8 +6,9 @@ # Third-Party Libraries from beanie import Document, Insert, Link, Replace, ValidateOnSave, before_event -from pydantic import BaseModel, EmailStr, Field, field_validator, ConfigDict +from pydantic import BaseModel, ConfigDict, EmailStr, Field, field_validator +from ..utils import utcnow from .enum import ( AgencyType, DayOfWeek, @@ -19,8 +20,6 @@ Stage, ) -from ..utils import utcnow - BOGUS_ID = "bogus_id_replace_me" diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index b1f239a..dbe9bc2 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -8,7 +8,7 @@ from beanie.operators import In, Push, Set from bson import ObjectId from bson.dbref import DBRef -from pydantic import Field, model_validator, ConfigDict +from pydantic import ConfigDict, Field, model_validator from pymongo import ASCENDING, IndexModel diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index be36fca..44f0343 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -1,15 +1,13 @@ # Standard Python Libraries from datetime import datetime from ipaddress import IPv4Network +from typing import Dict, List # Third-Party Libraries from beanie import Document -from pydantic import Field, ConfigDict -from pymongo import ASCENDING, IndexModel -from datetime import datetime -from pydantic import BaseModel, Field -from typing import List, Dict from bson import ObjectId +from pydantic import BaseModel, ConfigDict, Field +from pymongo import ASCENDING, IndexModel from ..utils import utcnow diff --git a/src/cyhy_db/models/system_control_doc.py b/src/cyhy_db/models/system_control_doc.py index 476125d..e59074b 100644 --- a/src/cyhy_db/models/system_control_doc.py +++ b/src/cyhy_db/models/system_control_doc.py @@ -1,11 +1,14 @@ -from beanie import Document +# Standard Python Libraries +import asyncio from datetime import datetime -from pydantic import Field, ConfigDict from typing import Optional -import asyncio -from .enum import ControlAction, ControlTarget +# Third-Party Libraries +from beanie import Document +from pydantic import ConfigDict, Field + from ..utils import utcnow +from .enum import ControlAction, ControlTarget CONTROL_DOC_POLL_INTERVAL = 5 # seconds diff --git a/src/cyhy_db/models/tally_doc.py b/src/cyhy_db/models/tally_doc.py index 5be3ce8..37c1a57 100644 --- a/src/cyhy_db/models/tally_doc.py +++ b/src/cyhy_db/models/tally_doc.py @@ -1,6 +1,9 @@ -from beanie import Document, before_event, Insert, Replace, ValidateOnSave +# Standard Python Libraries from datetime import datetime -from pydantic import BaseModel, Field, ConfigDict + +# Third-Party Libraries +from beanie import Document, Insert, Replace, ValidateOnSave, before_event +from pydantic import BaseModel, ConfigDict, Field from ..utils import utcnow diff --git a/src/cyhy_db/models/vuln_scan_doc.py b/src/cyhy_db/models/vuln_scan_doc.py index 65f31ff..50d2239 100644 --- a/src/cyhy_db/models/vuln_scan_doc.py +++ b/src/cyhy_db/models/vuln_scan_doc.py @@ -1,7 +1,10 @@ +# Standard Python Libraries +from datetime import datetime from typing import Dict + +# Third-Party Libraries from pydantic import ConfigDict from pymongo import ASCENDING, IndexModel -from datetime import datetime from . import ScanDoc from .enum import Protocol diff --git a/src/cyhy_db/utils/__init__.py b/src/cyhy_db/utils/__init__.py index 9cdfc87..72d8685 100644 --- a/src/cyhy_db/utils/__init__.py +++ b/src/cyhy_db/utils/__init__.py @@ -3,5 +3,4 @@ from .decorators import deprecated from .time import utcnow - __all__ = ["deprecated", "utcnow"] diff --git a/src/cyhy_db/utils/decorators.py b/src/cyhy_db/utils/decorators.py index ef3b8c8..83d20c1 100644 --- a/src/cyhy_db/utils/decorators.py +++ b/src/cyhy_db/utils/decorators.py @@ -1,3 +1,4 @@ +# Standard Python Libraries import warnings diff --git a/tests/test_data_generator.py b/tests/test_data_generator.py index 0a39d8d..4b538ac 100644 --- a/tests/test_data_generator.py +++ b/tests/test_data_generator.py @@ -1,9 +1,22 @@ +# Standard Python Libraries +from datetime import datetime +import ipaddress +import random + +# Third-Party Libraries +import factory +from mimesis import Generic +from mimesis.locales import DEFAULT_LOCALE +from mimesis.providers.base import BaseProvider +from mimesis_factory import MimesisField +from pytest_factoryboy import register + +# cisagov Libraries from cyhy_db.models import CVE, RequestDoc -from cyhy_db.models.request_doc import Agency, Contact, Location, Window from cyhy_db.models.enum import ( AgencyType, - DayOfWeek, CVSSVersion, + DayOfWeek, PocType, ReportPeriod, ReportType, @@ -11,16 +24,8 @@ Scheduler, Stage, ) +from cyhy_db.models.request_doc import Agency, Contact, Location, Window from cyhy_db.utils import utcnow -from mimesis_factory import MimesisField -from mimesis.locales import DEFAULT_LOCALE -from mimesis.providers.base import BaseProvider -from pytest_factoryboy import register -import factory -from mimesis import Generic -import ipaddress -import random -from datetime import datetime class CyHyProvider(BaseProvider): @@ -53,7 +58,7 @@ def network_ipv4(self): @register class CVEFactory(factory.Factory): - class Meta(object): + class Meta: model = CVE id = factory.LazyFunction(lambda: generic.cyhy_provider.cve_id()) @@ -63,7 +68,7 @@ class Meta(object): class AgencyFactory(factory.Factory): - class Meta(object): + class Meta: model = Agency name = factory.Faker("company") @@ -79,7 +84,7 @@ class Meta(object): class ContactFactory(factory.Factory): - class Meta(object): + class Meta: model = Contact email = factory.Faker("email") @@ -89,7 +94,7 @@ class Meta(object): class LocationFactory(factory.Factory): - class Meta(object): + class Meta: model = Location country_name = factory.Faker("country") @@ -104,7 +109,7 @@ class Meta(object): class WindowFactory(factory.Factory): - class Meta(object): + class Meta: model = Window day = factory.LazyFunction(lambda: random.choice(list(DayOfWeek))) @@ -113,7 +118,7 @@ class Meta(object): class RequestDocFactory(factory.Factory): - class Meta(object): + class Meta: model = RequestDoc id = factory.LazyAttribute( @@ -138,9 +143,9 @@ class Meta(object): ) # create a set of 1 to 3 random scan types from the ScanType enum scan_types = factory.LazyFunction( - lambda: set( - [random.choice(list(ScanType)) for _ in range(random.randint(1, 3))] - ) + lambda: { + random.choice(list(ScanType)) for _ in range(random.randint(1, 3)) + } ) diff --git a/tests/test_host_doc.py b/tests/test_host_doc.py index 413504a..0cbcc37 100644 --- a/tests/test_host_doc.py +++ b/tests/test_host_doc.py @@ -1,6 +1,9 @@ -from cyhy_db.models import HostDoc +# Standard Python Libraries from ipaddress import ip_address +# cisagov Libraries +from cyhy_db.models import HostDoc + VALID_IP_1_STR = "0.0.0.1" VALID_IP_2_STR = "0.0.0.2" VALID_IP_1_INT = int(ip_address(VALID_IP_1_STR)) From 30d3b080d9a5f39b0e12832d56068efb6b08f663 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 13:41:20 -0400 Subject: [PATCH 058/139] Add import statement for SnapshotDoc in ScanDoc --- src/cyhy_db/models/scan_doc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index dbe9bc2..1e90a8b 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -11,6 +11,8 @@ from pydantic import ConfigDict, Field, model_validator from pymongo import ASCENDING, IndexModel +from .snapshot_doc import SnapshotDoc + class ScanDoc(Document): # Validate on assignment so ip_int is recalculated as ip is set From 30f82e87e928fe3377bf24303bdfd4fbc02110c7 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 13:41:59 -0400 Subject: [PATCH 059/139] Apply changes made by linters --- tests/test_data_generator.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/test_data_generator.py b/tests/test_data_generator.py index 4b538ac..b26723e 100644 --- a/tests/test_data_generator.py +++ b/tests/test_data_generator.py @@ -143,9 +143,7 @@ class Meta: ) # create a set of 1 to 3 random scan types from the ScanType enum scan_types = factory.LazyFunction( - lambda: { - random.choice(list(ScanType)) for _ in range(random.randint(1, 3)) - } + lambda: {random.choice(list(ScanType)) for _ in range(random.randint(1, 3))} ) From 024e45d455dd5b15282c3b27136a1d9e2d4006d0 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 13:42:11 -0400 Subject: [PATCH 060/139] Remove unused index in RequestDoc model --- src/cyhy_db/models/request_doc.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 9add79e..38425fc 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -107,4 +107,3 @@ async def set_id_to_acronym(self): class Settings: # Beanie settings name = "requests" - indexes = [] From da7e5670ace6a55a5127308f374c0346b97a0970 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 13 Mar 2024 13:42:35 -0400 Subject: [PATCH 061/139] Add module docstring to disable isort breaking import order. --- src/cyhy_db/models/__init__.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py index 05748be..7e2c947 100644 --- a/src/cyhy_db/models/__init__.py +++ b/src/cyhy_db/models/__init__.py @@ -1,3 +1,11 @@ +"""This module contains the models for the CyHy database. + +# Imports are ordered to avoid a circular import. +# isort is disabled for this file as it will break the ordering. + +isort:skip_file +""" + # Scan documents (order matters) from .scan_doc import ScanDoc from .host_scan_doc import HostScanDoc From be547a6958d38b3072e6a786df0fd447ee15ca7b Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 22 Apr 2024 11:42:06 -0400 Subject: [PATCH 062/139] Clean up unused imports --- src/cyhy_db/models/cve.py | 2 +- src/cyhy_db/models/notification_doc.py | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 54ac75e..17a4958 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -2,7 +2,7 @@ from typing import Any, Dict # Third-Party Libraries -from beanie import Document, Indexed, ValidateOnSave, before_event +from beanie import Document, Indexed from pydantic import ConfigDict, Field, model_validator from .enum import CVSSVersion diff --git a/src/cyhy_db/models/notification_doc.py b/src/cyhy_db/models/notification_doc.py index 922e71f..fdb87c1 100644 --- a/src/cyhy_db/models/notification_doc.py +++ b/src/cyhy_db/models/notification_doc.py @@ -3,8 +3,7 @@ # Third-Party Libraries from beanie import BeanieObjectId, Document -from bson import ObjectId -from pydantic import BaseModel, ConfigDict, Field +from pydantic import ConfigDict, Field class NotificationDoc(Document): From 36b2a220d80155bffebb9ec25beef1a31abffc11 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 22 Apr 2024 11:42:43 -0400 Subject: [PATCH 063/139] Add TODO to delete BASESCAN if unused --- src/cyhy_db/models/enum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index 29f749b..a8f157d 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -72,7 +72,7 @@ class Scheduler(Enum): class Stage(Enum): - BASESCAN = "BASESCAN" + BASESCAN = "BASESCAN" # TODO: Delete if unused NETSCAN1 = "NETSCAN1" NETSCAN2 = "NETSCAN2" PORTSCAN = "PORTSCAN" From a2b798af25534264564709e4bf45a90f66c2491c Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 22 Apr 2024 11:43:50 -0400 Subject: [PATCH 064/139] Improve comments on utcnow util --- src/cyhy_db/utils/time.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/utils/time.py b/src/cyhy_db/utils/time.py index c3e4512..c20a25b 100644 --- a/src/cyhy_db/utils/time.py +++ b/src/cyhy_db/utils/time.py @@ -3,5 +3,11 @@ def utcnow() -> datetime: - """Returns a timezone-aware datetime object with the current time in UTC.""" + """Returns a timezone-aware datetime object with the current time in UTC. + + This is useful for default value factories in Beanie models. + + Returns: + datetime: The current time in UTC + """ return datetime.now(timezone.utc) From b7c28b0020c339b8b5043ab97c8d40f4ab7432d0 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 22 Apr 2024 15:16:08 -0400 Subject: [PATCH 065/139] Replace date factory with custom function. --- src/cyhy_db/models/scan_doc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 1e90a8b..149241e 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -12,6 +12,7 @@ from pymongo import ASCENDING, IndexModel from .snapshot_doc import SnapshotDoc +from ..utils import utcnow class ScanDoc(Document): @@ -24,7 +25,7 @@ class ScanDoc(Document): owner: str = Field(...) snapshots: List[Link["SnapshotDoc"]] = Field(default=[]) source: str = Field(...) - time: datetime = Field(default_factory=datetime.utcnow) + time: datetime = Field(default_factory=utcnow) @model_validator(mode="before") def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: From c16c4e47e346ea09c5b1d067bd0f0d614b26d8ca Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 22 Apr 2024 15:16:59 -0400 Subject: [PATCH 066/139] Add ScanDoc model header comment --- src/cyhy_db/models/scan_doc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 149241e..17fa6fa 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -1,3 +1,5 @@ +""" ScanDoc model for use as the base of other scan document classes. """ + # Standard Python Libraries from datetime import datetime from ipaddress import IPv4Address, ip_address From d44ad805a53801c1be0c071d002c131c7d96d886 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 22 Apr 2024 15:17:49 -0400 Subject: [PATCH 067/139] Add TODO to inherit from BaseModel I can't remember why I added this TODO. Someday I will, and that day will be glorious. --- src/cyhy_db/models/scan_doc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 17fa6fa..d126500 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -10,14 +10,14 @@ from beanie.operators import In, Push, Set from bson import ObjectId from bson.dbref import DBRef -from pydantic import ConfigDict, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from pymongo import ASCENDING, IndexModel from .snapshot_doc import SnapshotDoc from ..utils import utcnow -class ScanDoc(Document): +class ScanDoc(Document): # TODO: Make this a BaseModel # Validate on assignment so ip_int is recalculated as ip is set model_config = ConfigDict(extra="forbid", validate_assignment=True) From c81cc5a913babb4571118e0fb9a86bae39beda46 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 26 Aug 2024 14:53:33 -0400 Subject: [PATCH 068/139] Add non-interactive flag to mypy types install The script will pause and ask for user input without this flag. --- setup-env | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup-env b/setup-env index ac7ecfc..b3554cb 100755 --- a/setup-env +++ b/setup-env @@ -251,7 +251,7 @@ for req_file in "requirements-dev.txt" "requirements-test.txt" "requirements.txt done # Install all necessary mypy type stubs -mypy --install-types src/ +mypy --install-types --non-interactive src/ # Install git pre-commit hooks now or later. pre-commit install ${INSTALL_HOOKS:+"--install-hooks"} From 0d9d1631ed725a4cc492d597fda63e138c9285cd Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 26 Aug 2024 15:15:46 -0400 Subject: [PATCH 069/139] Bump docker library to fix upstream break by requests See: - https://github.com/docker/docker-py/pull/3257 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 589ac3f..61e3542 100644 --- a/setup.py +++ b/setup.py @@ -108,7 +108,7 @@ def get_version(version_file): # 1.11.1 fixed this issue, but to ensure expected behavior we'll pin # to never grab the regression version. "coveralls != 1.11.0", - "docker == 7.0.0", + "docker == 7.1.0", "hypothesis", "mimesis-factory", "mimesis", From 8b9fd65a8904d17d8903684f1e76bce982bed434 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 26 Aug 2024 16:53:19 -0400 Subject: [PATCH 070/139] Add types and ignores to make mypy happy --- src/cyhy_db/db.py | 12 ++++++------ src/cyhy_db/models/cve.py | 3 ++- src/cyhy_db/models/host_doc.py | 3 ++- src/cyhy_db/models/kev_doc.py | 4 ++-- src/cyhy_db/models/place_doc.py | 5 +++-- src/cyhy_db/models/request_doc.py | 12 ++++++++++-- src/cyhy_db/models/system_control_doc.py | 2 +- 7 files changed, 26 insertions(+), 15 deletions(-) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index cfb797a..cc49bbb 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -1,10 +1,10 @@ # Third-Party Libraries -from beanie import init_beanie -from motor.motor_asyncio import AsyncIOMotorClient +from beanie import Document, View, init_beanie +from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase from .models import * -ALL_MODELS = [ +ALL_MODELS: list[type[Document] | type[View] | str] = [ CVE, HostDoc, HostScanDoc, @@ -22,10 +22,10 @@ ] -async def initialize_db(db_uri: str, db_name: str) -> None: +async def initialize_db(db_uri: str, db_name: str) -> AsyncIOMotorDatabase: try: - client = AsyncIOMotorClient(db_uri) - db = client[db_name] + client: AsyncIOMotorClient = AsyncIOMotorClient(db_uri) + db: AsyncIOMotorDatabase = client[db_name] await init_beanie(database=db, document_models=ALL_MODELS) return db except Exception as e: diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 17a4958..a37e9b2 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -12,7 +12,8 @@ class CVE(Document): # Validate on assignment so ip_int is recalculated as ip is set model_config = ConfigDict(extra="forbid", validate_assignment=True) - id: str = Indexed(primary_field=True) # CVE ID + # CVE ID as a string + id: str = Indexed(primary_field=True) # type: ignore[assignment] cvss_score: float = Field(ge=0.0, le=10.0) cvss_version: CVSSVersion = Field(default=CVSSVersion.V3_1) severity: int = Field(ge=1, le=4, default=1) diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 3339b70..d612516 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -21,7 +21,8 @@ class State(BaseModel): class HostDoc(Document): model_config = ConfigDict(extra="forbid") - id: int = Field() # IP address as an integer + # IP address as an integer + id: int = Field(default_factory=int) # type: ignore[assignment] ip: IPv4Address = Field(...) owner: str = Field(...) last_change: datetime = Field(default_factory=utcnow) diff --git a/src/cyhy_db/models/kev_doc.py b/src/cyhy_db/models/kev_doc.py index 20859bf..240649a 100644 --- a/src/cyhy_db/models/kev_doc.py +++ b/src/cyhy_db/models/kev_doc.py @@ -1,12 +1,12 @@ # Third-Party Libraries from beanie import Document -from pydantic import ConfigDict +from pydantic import ConfigDict, Field class KEVDoc(Document): model_config = ConfigDict(extra="forbid") - id: str # CVE + id: str = Field(default_factory=str) # type: ignore[assignment] known_ransomware: bool class Settings: diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index a0ad557..aefd7fb 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -2,14 +2,15 @@ from typing import Optional # Third-Party Libraries -from beanie import Document +from beanie import Document, PydanticObjectId from pydantic import ConfigDict, Field class PlaceDoc(Document): model_config = ConfigDict(extra="forbid") - id: int # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html + # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html + id: int = Field(default_factory=int) # type: ignore[assignment] name: str clazz: str = Field(alias="class") # 'class' is a reserved keyword in Python state: str diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 38425fc..78e80db 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -5,7 +5,15 @@ from typing import List, Optional # Third-Party Libraries -from beanie import Document, Insert, Link, Replace, ValidateOnSave, before_event +from beanie import ( + Document, + Insert, + Link, + PydanticObjectId, + Replace, + ValidateOnSave, + before_event, +) from pydantic import BaseModel, ConfigDict, EmailStr, Field, field_validator from ..utils import utcnow @@ -81,7 +89,7 @@ def validate_start(cls, v): class RequestDoc(Document): model_config = ConfigDict(extra="forbid") - id: str = Field(default=BOGUS_ID) + id: str = Field(default=BOGUS_ID) # type: ignore[assignment] agency: Agency children: List[Link["RequestDoc"]] = Field(default=[]) enrolled: datetime = Field(default_factory=utcnow) diff --git a/src/cyhy_db/models/system_control_doc.py b/src/cyhy_db/models/system_control_doc.py index e59074b..aba46ce 100644 --- a/src/cyhy_db/models/system_control_doc.py +++ b/src/cyhy_db/models/system_control_doc.py @@ -33,7 +33,7 @@ async def wait_for_completion(cls, document_id, timeout: Optional[int] = None): start_time = utcnow() while True: doc = await cls.get(document_id) - if doc.completed: + if doc and doc.completed: return True if timeout and (utcnow() - start_time).total_seconds() > timeout: return False From a33446cc745b18f3363e323ce921263b2f9177d6 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 27 Aug 2024 10:48:46 -0400 Subject: [PATCH 071/139] Resolve some flake8 issues * Remove unused imports * Specify all imports, instead of using wildcard imports * Use preferred comparisons * Adhere to docstring whitespace and imperative mood rules --- src/cyhy_db/db.py | 17 ++++++++++++++++- src/cyhy_db/models/host_doc.py | 18 +++++++++--------- src/cyhy_db/models/place_doc.py | 2 +- src/cyhy_db/models/report_doc.py | 2 +- src/cyhy_db/models/request_doc.py | 1 - src/cyhy_db/models/scan_doc.py | 6 +++--- src/cyhy_db/models/snapshot_doc.py | 1 - src/cyhy_db/models/system_control_doc.py | 7 +++++-- src/cyhy_db/models/vuln_scan_doc.py | 1 - src/cyhy_db/utils/time.py | 2 +- tests/conftest.py | 3 --- tests/test_connection.py | 2 +- tests/test_data_generator.py | 2 -- tests/test_request_doc.py | 9 --------- tests/test_scan_doc.py | 13 ++++++------- 15 files changed, 43 insertions(+), 43 deletions(-) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index cc49bbb..9346f09 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -2,7 +2,22 @@ from beanie import Document, View, init_beanie from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase -from .models import * +from .models import ( + CVE, + HostDoc, + HostScanDoc, + KEVDoc, + NotificationDoc, + PlaceDoc, + PortScanDoc, + RequestDoc, + ReportDoc, + ScanDoc, + SnapshotDoc, + SystemControlDoc, + TallyDoc, + VulnScanDoc, +) ALL_MODELS: list[type[Document] | type[View] | str] = [ CVE, diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index d612516..b55f396 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -5,7 +5,7 @@ from typing import Any, Dict, Optional, Tuple # Third-Party Libraries -from beanie import Document, Indexed, Insert, Replace, ValidateOnSave, before_event +from beanie import Document, Insert, Replace, ValidateOnSave, before_event from pydantic import BaseModel, ConfigDict, Field, model_validator from pymongo import ASCENDING, IndexModel @@ -98,17 +98,17 @@ class Settings: ] def set_state(self, nmap_says_up, has_open_ports, reason=None): - """Sets state.up based on different stage - evidence. nmap has a concept of up which is - different from our definition. An nmap "up" just - means it got a reply, not that there are any open - ports. Note either argument can be None.""" + """Set state.up based on different stage evidence. - if has_open_ports == True: # Only PORTSCAN sends in has_open_ports + nmap has a concept of up which is different from our definition. An nmap + "up" just means it got a reply, not that there are any open ports. Note + either argument can be None. + """ + if has_open_ports: # Only PORTSCAN sends in has_open_ports self.state = State(True, "open-port") - elif has_open_ports == False: + elif has_open_ports is False: self.state = State(False, "no-open") - elif nmap_says_up == False: # NETSCAN says host is down + elif nmap_says_up is False: # NETSCAN says host is down self.state = State(False, reason) # TODO: There are a lot of functions in the Python 2 version that may or may not be used. diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index aefd7fb..b0aa022 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -2,7 +2,7 @@ from typing import Optional # Third-Party Libraries -from beanie import Document, PydanticObjectId +from beanie import Document from pydantic import ConfigDict, Field diff --git a/src/cyhy_db/models/report_doc.py b/src/cyhy_db/models/report_doc.py index b917f80..62c89d3 100644 --- a/src/cyhy_db/models/report_doc.py +++ b/src/cyhy_db/models/report_doc.py @@ -3,7 +3,7 @@ from typing import List # Third-Party Libraries -from beanie import Document, Indexed, Link +from beanie import Document, Link from pydantic import ConfigDict, Field from pymongo import ASCENDING, IndexModel diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 78e80db..9c74747 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -9,7 +9,6 @@ Document, Insert, Link, - PydanticObjectId, Replace, ValidateOnSave, before_event, diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index d126500..38dcded 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -1,4 +1,4 @@ -""" ScanDoc model for use as the base of other scan document classes. """ +"""ScanDoc model for use as the base of other scan document classes.""" # Standard Python Libraries from datetime import datetime @@ -10,11 +10,11 @@ from beanie.operators import In, Push, Set from bson import ObjectId from bson.dbref import DBRef -from pydantic import BaseModel, ConfigDict, Field, model_validator +from pydantic import ConfigDict, Field, model_validator from pymongo import ASCENDING, IndexModel -from .snapshot_doc import SnapshotDoc from ..utils import utcnow +from .snapshot_doc import SnapshotDoc class ScanDoc(Document): # TODO: Make this a BaseModel diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index 44f0343..4f76d0b 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -5,7 +5,6 @@ # Third-Party Libraries from beanie import Document -from bson import ObjectId from pydantic import BaseModel, ConfigDict, Field from pymongo import ASCENDING, IndexModel diff --git a/src/cyhy_db/models/system_control_doc.py b/src/cyhy_db/models/system_control_doc.py index aba46ce..606f575 100644 --- a/src/cyhy_db/models/system_control_doc.py +++ b/src/cyhy_db/models/system_control_doc.py @@ -28,8 +28,11 @@ class Settings: @classmethod async def wait_for_completion(cls, document_id, timeout: Optional[int] = None): - """Wait for this control action to complete. If a timeout is set, only wait a maximum of timeout seconds. - Returns True if the document was completed, False otherwise.""" + """Wait for this control action to complete. + + If a timeout is set, only wait a maximum of timeout seconds. + Returns True if the document was completed, False otherwise. + """ start_time = utcnow() while True: doc = await cls.get(document_id) diff --git a/src/cyhy_db/models/vuln_scan_doc.py b/src/cyhy_db/models/vuln_scan_doc.py index 50d2239..0b6ac6c 100644 --- a/src/cyhy_db/models/vuln_scan_doc.py +++ b/src/cyhy_db/models/vuln_scan_doc.py @@ -1,6 +1,5 @@ # Standard Python Libraries from datetime import datetime -from typing import Dict # Third-Party Libraries from pydantic import ConfigDict diff --git a/src/cyhy_db/utils/time.py b/src/cyhy_db/utils/time.py index c20a25b..683e660 100644 --- a/src/cyhy_db/utils/time.py +++ b/src/cyhy_db/utils/time.py @@ -3,7 +3,7 @@ def utcnow() -> datetime: - """Returns a timezone-aware datetime object with the current time in UTC. + """Return a timezone-aware datetime object with the current time in UTC. This is useful for default value factories in Beanie models. diff --git a/tests/conftest.py b/tests/conftest.py index 561d280..be2c53e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,12 +9,9 @@ import time # Third-Party Libraries -from beanie import Document, init_beanie import docker from motor.core import AgnosticClient -from motor.motor_asyncio import AsyncIOMotorClient import pytest -from pytest_asyncio import is_async_test # cisagov Libraries from cyhy_db import initialize_db diff --git a/tests/test_connection.py b/tests/test_connection.py index a26f9fb..0fe7564 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -4,7 +4,7 @@ from motor.motor_asyncio import AsyncIOMotorClient # cisagov Libraries -from cyhy_db import initialize_db +# from cyhy_db import initialize_db from cyhy_db.models import CVE diff --git a/tests/test_data_generator.py b/tests/test_data_generator.py index b26723e..2ee0bb5 100644 --- a/tests/test_data_generator.py +++ b/tests/test_data_generator.py @@ -8,7 +8,6 @@ from mimesis import Generic from mimesis.locales import DEFAULT_LOCALE from mimesis.providers.base import BaseProvider -from mimesis_factory import MimesisField from pytest_factoryboy import register # cisagov Libraries @@ -19,7 +18,6 @@ DayOfWeek, PocType, ReportPeriod, - ReportType, ScanType, Scheduler, Stage, diff --git a/tests/test_request_doc.py b/tests/test_request_doc.py index d4e627d..95d511d 100644 --- a/tests/test_request_doc.py +++ b/tests/test_request_doc.py @@ -1,14 +1,5 @@ """Test RequestDoc model functionality.""" -# Standard Python Libraries -import ipaddress - -# Third-Party Libraries -from hypothesis import given -from hypothesis import strategies as st -from pydantic import ValidationError -import pytest - # cisagov Libraries from cyhy_db.models import RequestDoc from cyhy_db.models.request_doc import Agency diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index a24676f..6ac4531 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -1,7 +1,6 @@ """Test ScanDoc model functionality.""" # Standard Python Libraries -import datetime import ipaddress # Third-Party Libraries @@ -53,7 +52,7 @@ def test_ip_string_set(): source="nmap", ) - assert type(scan_doc.ip) == ipaddress.IPv4Address, "IP address was not converted" + assert isinstance(scan_doc.ip, ipaddress.IPv4Address), "IP address was not converted" assert scan_doc.ip_int == VALID_IP_1_INT, "IP address integer was not calculated" @@ -96,13 +95,13 @@ async def test_reset_latest_flag_by_owner(): ) await scan_doc.save() # Check that the latest flag is set to True - assert scan_doc.latest == True + assert scan_doc.latest is True # Reset the latest flag await ScanDoc.reset_latest_flag_by_owner(OWNER) # Retrieve the ScanDoc object from the database await scan_doc.sync() # Check that the latest flag is set to False - assert scan_doc.latest == False + assert scan_doc.latest is False async def test_reset_latest_flag_by_ip(): @@ -114,19 +113,19 @@ async def test_reset_latest_flag_by_ip(): await scan_doc_1.save() await scan_doc_2.save() # Check that the latest flag is set to True - assert scan_doc_1.latest == True + assert scan_doc_1.latest is True # Reset the latest flag on single IP await ScanDoc.reset_latest_flag_by_ip(IP_TO_RESET_1) # Retrieve the ScanDoc object from the database await scan_doc_1.sync() # Check that the latest flag is set to False - assert scan_doc_1.latest == False + assert scan_doc_1.latest is False # Reset by both IPs await ScanDoc.reset_latest_flag_by_ip([IP_TO_RESET_1, IP_TO_RESET_2]) # Retrieve the ScanDoc object from the database await scan_doc_2.sync() # Check that the latest flag is set to False - assert scan_doc_2.latest == False + assert scan_doc_2.latest is False async def test_tag_latest(): From 4d6050474201dadbf41d35e1070df3c37238da34 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 27 Aug 2024 10:48:45 -0400 Subject: [PATCH 072/139] Bump dependency versions to latest --- setup.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/setup.py b/setup.py index 61e3542..91f2097 100644 --- a/setup.py +++ b/setup.py @@ -91,11 +91,11 @@ def get_version(version_file): py_modules=[splitext(basename(path))[0] for path in glob("src/*.py")], include_package_data=True, install_requires=[ - "beanie == 1.25.0", + "beanie == 1.26.0", "docopt == 0.6.2", "pydantic[email, hypothesis]", # hypothesis plugin is currently disabled: https://github.com/pydantic/pydantic/issues/4682 - "schema == 0.7.5", - "setuptools >= 69.0.3", + "schema == 0.7.7", + "setuptools >= 73.0.1", ], extras_require={ "test": [ From 1b5be0087a375d293f85e33dd78fba489745aa75 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 27 Aug 2024 10:56:18 -0400 Subject: [PATCH 073/139] Satisfy isort and black linters --- src/cyhy_db/db.py | 2 +- src/cyhy_db/models/request_doc.py | 9 +-------- tests/test_scan_doc.py | 4 +++- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index 9346f09..93f7d7f 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -10,8 +10,8 @@ NotificationDoc, PlaceDoc, PortScanDoc, - RequestDoc, ReportDoc, + RequestDoc, ScanDoc, SnapshotDoc, SystemControlDoc, diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 9c74747..cec041a 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -5,14 +5,7 @@ from typing import List, Optional # Third-Party Libraries -from beanie import ( - Document, - Insert, - Link, - Replace, - ValidateOnSave, - before_event, -) +from beanie import Document, Insert, Link, Replace, ValidateOnSave, before_event from pydantic import BaseModel, ConfigDict, EmailStr, Field, field_validator from ..utils import utcnow diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index 6ac4531..682c375 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -52,7 +52,9 @@ def test_ip_string_set(): source="nmap", ) - assert isinstance(scan_doc.ip, ipaddress.IPv4Address), "IP address was not converted" + assert isinstance( + scan_doc.ip, ipaddress.IPv4Address + ), "IP address was not converted" assert scan_doc.ip_int == VALID_IP_1_INT, "IP address integer was not calculated" From 8fdcfda446e92eb49fbdf3ade6267d4a0968ff44 Mon Sep 17 00:00:00 2001 From: Felddy Date: Thu, 29 Aug 2024 15:56:30 -0400 Subject: [PATCH 074/139] Update comment --- src/cyhy_db/models/cve.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index a37e9b2..03858eb 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -9,7 +9,7 @@ class CVE(Document): - # Validate on assignment so ip_int is recalculated as ip is set + # Validate on assignment so severity is calculated model_config = ConfigDict(extra="forbid", validate_assignment=True) # CVE ID as a string From 5f8b975dec42f79171349e41230cc9bac6b3bf0a Mon Sep 17 00:00:00 2001 From: Felddy Date: Thu, 29 Aug 2024 15:57:48 -0400 Subject: [PATCH 075/139] Add py.typed file to allow users to type check --- setup.py | 2 +- src/cyhy_db/py.typed | 0 2 files changed, 1 insertion(+), 1 deletion(-) create mode 100644 src/cyhy_db/py.typed diff --git a/setup.py b/setup.py index 91f2097..edde11e 100644 --- a/setup.py +++ b/setup.py @@ -87,7 +87,7 @@ def get_version(version_file): keywords=["cyhy", "database"], packages=find_packages(where="src"), package_dir={"": "src"}, - package_data={}, + package_data={"cyhy_db": ["py.typed"]}, py_modules=[splitext(basename(path))[0] for path in glob("src/*.py")], include_package_data=True, install_requires=[ diff --git a/src/cyhy_db/py.typed b/src/cyhy_db/py.typed new file mode 100644 index 0000000..e69de29 From 59e95950def4e90a7a2a5cb3778285affe525ee8 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Fri, 30 Aug 2024 11:00:40 -0400 Subject: [PATCH 076/139] Ignore some E712 flake8 findings Beanie syntax requires us to write code in a way that triggers E712. --- src/cyhy_db/models/scan_doc.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 38dcded..794401c 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -49,7 +49,9 @@ class Settings: @classmethod async def reset_latest_flag_by_owner(cls, owner: str): - await cls.find(cls.latest == True, cls.owner == owner).update_many( + # flake8 E712 is "comparison to True should be 'if cond is True:' or 'if + # cond:'" but this is unavoidable due to Beanie syntax. + await cls.find(cls.latest == True, cls.owner == owner).update_many( # noqa E712 Set({cls.latest: False}) ) @@ -71,9 +73,11 @@ async def reset_latest_flag_by_ip( else: ip_ints = [int(ip_address(ips))] - await cls.find(cls.latest == True, In(cls.ip_int, ip_ints)).update_many( - Set({cls.latest: False}) - ) + # flake8 E712 is "comparison to True should be 'if cond is True:' or 'if + # cond:'" but this is unavoidable due to Beanie syntax. + await cls.find( + cls.latest == True, In(cls.ip_int, ip_ints) # noqa E712 + ).update_many(Set({cls.latest: False})) @classmethod async def tag_latest( @@ -89,6 +93,8 @@ async def tag_latest( ref = DBRef(SnapshotDoc.Settings.name, ObjectId(snapshot)) else: raise ValueError("Invalid snapshot type") - await cls.find(cls.latest == True, In(cls.owner, owners)).update_many( - Push({cls.snapshots: ref}) - ) + # flake8 E712 is "comparison to True should be 'if cond is True:' or 'if + # cond:'" but this is unavoidable due to Beanie syntax. + await cls.find( + cls.latest == True, In(cls.owner, owners) # noqa E712 + ).update_many(Push({cls.snapshots: ref})) From d1489e483a646589ecbf61d4c91467c3a6b7ea82 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Fri, 30 Aug 2024 11:02:05 -0400 Subject: [PATCH 077/139] Appease flake8 by using an otherwise unused variable Error was: F841 local variable 'host_doc' is assigned to but never used --- tests/test_host_doc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_host_doc.py b/tests/test_host_doc.py index 0cbcc37..73130bf 100644 --- a/tests/test_host_doc.py +++ b/tests/test_host_doc.py @@ -16,6 +16,8 @@ def test_host_doc_init(): ip=ip_address(VALID_IP_1_STR), owner="YOUR_MOM", ) + # Check that the HostDoc object was created correctly + assert host_doc.ip == ip_address(VALID_IP_1_STR) async def test_save(): From 757caa85fec05bb973726c8c4640496693b8cf8a Mon Sep 17 00:00:00 2001 From: David Redmin Date: Fri, 30 Aug 2024 15:48:37 -0400 Subject: [PATCH 078/139] Add a whole bunch of docstrings to placate flake8 --- src/cyhy_db/__init__.py | 2 + src/cyhy_db/db.py | 3 ++ src/cyhy_db/models/cve.py | 8 ++- src/cyhy_db/models/enum.py | 30 +++++++++++ src/cyhy_db/models/host_doc.py | 12 ++++- src/cyhy_db/models/host_scan_doc.py | 7 ++- src/cyhy_db/models/kev_doc.py | 6 +++ src/cyhy_db/models/notification_doc.py | 6 +++ src/cyhy_db/models/place_doc.py | 6 +++ src/cyhy_db/models/port_scan_doc.py | 7 ++- src/cyhy_db/models/report_doc.py | 6 +++ src/cyhy_db/models/request_doc.py | 21 ++++++-- src/cyhy_db/models/scan_doc.py | 9 +++- src/cyhy_db/models/snapshot_doc.py | 17 ++++++- src/cyhy_db/models/system_control_doc.py | 6 +++ src/cyhy_db/models/tally_doc.py | 11 +++++ src/cyhy_db/models/vuln_scan_doc.py | 7 ++- src/cyhy_db/utils/decorators.py | 4 ++ src/cyhy_db/utils/time.py | 2 + tests/conftest.py | 1 + tests/test_connection.py | 2 + tests/test_data_generator.py | 63 ++++++++++++++++++++---- tests/test_host_doc.py | 5 ++ tests/test_request_doc.py | 2 +- tests/test_scan_doc.py | 43 +++++++++++++++- 25 files changed, 264 insertions(+), 22 deletions(-) diff --git a/src/cyhy_db/__init__.py b/src/cyhy_db/__init__.py index 660c2b2..4b9e59d 100644 --- a/src/cyhy_db/__init__.py +++ b/src/cyhy_db/__init__.py @@ -1,3 +1,5 @@ +"""The cyhy_db package provides an interface to a CyHy database.""" + from .db import initialize_db __all__ = ["initialize_db"] diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index 93f7d7f..6213484 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -1,3 +1,5 @@ +"""CyHy database top-level functions.""" + # Third-Party Libraries from beanie import Document, View, init_beanie from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase @@ -38,6 +40,7 @@ async def initialize_db(db_uri: str, db_name: str) -> AsyncIOMotorDatabase: + """Initialize the database.""" try: client: AsyncIOMotorClient = AsyncIOMotorClient(db_uri) db: AsyncIOMotorDatabase = client[db_name] diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve.py index 03858eb..1c65370 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve.py @@ -1,3 +1,5 @@ +"""The model for CVE (Common Vulnerabilities and Exposures) documents.""" + # Standard Python Libraries from typing import Any, Dict @@ -9,6 +11,8 @@ class CVE(Document): + """The CVE document model.""" + # Validate on assignment so severity is calculated model_config = ConfigDict(extra="forbid", validate_assignment=True) @@ -20,6 +24,7 @@ class CVE(Document): @model_validator(mode="before") def calculate_severity(cls, values: Dict[str, Any]) -> Dict[str, Any]: + """Calculate CVE severity based on the CVSS score and version.""" if values["cvss_version"] == "2.0": if values["cvss_score"] == 10: values["severity"] = 4 @@ -41,5 +46,6 @@ def calculate_severity(cls, values: Dict[str, Any]) -> Dict[str, Any]: return values class Settings: - # Beanie settings + """Beanie settings.""" + name = "cves" diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index a8f157d..270d9bb 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -1,8 +1,12 @@ +"""The enumerations used in CyHy.""" + # Standard Python Libraries from enum import Enum class AgencyType(Enum): + """Agency types.""" + FEDERAL = "FEDERAL" LOCAL = "LOCAL" PRIVATE = "PRIVATE" @@ -12,21 +16,29 @@ class AgencyType(Enum): class ControlAction(Enum): + """Commander control actions.""" + PAUSE = "PAUSE" STOP = "STOP" class ControlTarget(Enum): + """Commander control targets.""" + COMMANDER = "COMMANDER" class CVSSVersion(Enum): + """CVSS versions.""" + V2 = "2.0" V3 = "3.0" V3_1 = "3.1" class DayOfWeek(Enum): + """Days of the week.""" + MONDAY = "MONDAY" TUESDAY = "TUESDAY" WEDNESDAY = "WEDNESDAY" @@ -37,22 +49,30 @@ class DayOfWeek(Enum): class PocType(Enum): + """Point of contact types.""" + DISTRO = "DISTRO" TECHNICAL = "TECHNICAL" class Protocol(Enum): + """Protocols.""" + TCP = "tcp" UDP = "udp" class ReportPeriod(Enum): + """CyHy reporting periods.""" + MONTHLY = "MONTHLY" QUARTERLY = "QUARTERLY" WEEKLY = "WEEKLY" class ReportType(Enum): + """CyHy report types.""" + BOD = "BOD" CYBEX = "CYBEX" CYHY = "CYHY" @@ -62,16 +82,22 @@ class ReportType(Enum): class ScanType(Enum): + """CyHy scan types.""" + CYHY = "CYHY" DNSSEC = "DNSSEC" PHISHING = "PHISHING" class Scheduler(Enum): + """CyHy schedulers.""" + PERSISTENT1 = "PERSISTENT1" class Stage(Enum): + """CyHy scan stages.""" + BASESCAN = "BASESCAN" # TODO: Delete if unused NETSCAN1 = "NETSCAN1" NETSCAN2 = "NETSCAN2" @@ -80,6 +106,8 @@ class Stage(Enum): class Status(Enum): + """CyHy scan statuses.""" + DONE = "DONE" READY = "READY" RUNNING = "RUNNING" @@ -87,6 +115,8 @@ class Status(Enum): class TicketEvent(Enum): + """Ticket events.""" + CHANGED = "CHANGED" CLOSED = "CLOSED" OPENED = "OPENED" diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index b55f396..1f28eeb 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy host documents.""" + # Standard Python Libraries from datetime import datetime from ipaddress import IPv4Address, ip_address @@ -14,11 +16,15 @@ class State(BaseModel): + """The state of a host.""" + reason: str up: bool class HostDoc(Document): + """The host document model.""" + model_config = ConfigDict(extra="forbid") # IP address as an integer @@ -37,16 +43,19 @@ class HostDoc(Document): @model_validator(mode="before") def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: + """Calculate the integer representation of an IP address.""" # ip may still be string if it was just set values["_id"] = int(ip_address(values["ip"])) return values @before_event(Insert, Replace, ValidateOnSave) async def before_save(self): + """Set data just prior to saving a host document.""" self.last_change = utcnow() class Settings: - # Beanie settings + """Beanie settings.""" + name = "hosts" indexes = [ IndexModel( @@ -124,6 +133,7 @@ def set_state(self, nmap_says_up, has_open_ports, reason=None): @classmethod @deprecated("Use HostDoc.find_one(HostDoc.ip == ip) instead.") async def get_by_ip(cls, ip: IPv4Address): + """Return a host document with the given IP address.""" return await cls.find_one(cls.ip == ip) # @classmethod diff --git a/src/cyhy_db/models/host_scan_doc.py b/src/cyhy_db/models/host_scan_doc.py index 010f350..ea2a1ad 100644 --- a/src/cyhy_db/models/host_scan_doc.py +++ b/src/cyhy_db/models/host_scan_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy host scan documents.""" + # Standard Python Libraries from typing import List @@ -9,6 +11,8 @@ class HostScanDoc(ScanDoc): + """The host scan document model.""" + model_config = ConfigDict(extra="forbid") name: str @@ -17,7 +21,8 @@ class HostScanDoc(ScanDoc): classes: List[dict] = [] class Settings: - # Beanie settings + """Beanie settings.""" + name = "host_scans" indexes = [ IndexModel( diff --git a/src/cyhy_db/models/kev_doc.py b/src/cyhy_db/models/kev_doc.py index 240649a..4f9986d 100644 --- a/src/cyhy_db/models/kev_doc.py +++ b/src/cyhy_db/models/kev_doc.py @@ -1,13 +1,19 @@ +"""The model for KEV (Known Exploited Vulnerabilities) documents.""" + # Third-Party Libraries from beanie import Document from pydantic import ConfigDict, Field class KEVDoc(Document): + """The KEV document model.""" + model_config = ConfigDict(extra="forbid") id: str = Field(default_factory=str) # type: ignore[assignment] known_ransomware: bool class Settings: + """Beanie settings.""" + name = "kevs" diff --git a/src/cyhy_db/models/notification_doc.py b/src/cyhy_db/models/notification_doc.py index fdb87c1..1261395 100644 --- a/src/cyhy_db/models/notification_doc.py +++ b/src/cyhy_db/models/notification_doc.py @@ -1,3 +1,5 @@ +"""The model for notification documents.""" + # Standard Python Libraries from typing import List @@ -7,6 +9,8 @@ class NotificationDoc(Document): + """The notification document model.""" + model_config = ConfigDict(extra="forbid") ticket_id: BeanieObjectId = Field(...) # ticket id that triggered the notification @@ -16,4 +20,6 @@ class NotificationDoc(Document): ) # list of owners built as notifications are generated class Settings: + """Beanie settings.""" + name = "notifications" diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index b0aa022..e220a85 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -1,3 +1,5 @@ +"""The model for place documents.""" + # Standard Python Libraries from typing import Optional @@ -7,6 +9,8 @@ class PlaceDoc(Document): + """The place document model.""" + model_config = ConfigDict(extra="forbid") # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html @@ -28,4 +32,6 @@ class PlaceDoc(Document): elevation_feet: Optional[int] = None class Settings: + """Beanie settings.""" + name = "places" diff --git a/src/cyhy_db/models/port_scan_doc.py b/src/cyhy_db/models/port_scan_doc.py index 9af2801..2236f1b 100644 --- a/src/cyhy_db/models/port_scan_doc.py +++ b/src/cyhy_db/models/port_scan_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy port scan documents.""" + # Standard Python Libraries from typing import Dict @@ -10,6 +12,8 @@ class PortScanDoc(ScanDoc): + """The port scan document model.""" + model_config = ConfigDict(extra="forbid") protocol: Protocol port: int @@ -18,7 +22,8 @@ class PortScanDoc(ScanDoc): reason: str class Settings: - # Beanie settings + """Beanie settings.""" + name = "port_scans" indexes = [ IndexModel( diff --git a/src/cyhy_db/models/report_doc.py b/src/cyhy_db/models/report_doc.py index 62c89d3..4526707 100644 --- a/src/cyhy_db/models/report_doc.py +++ b/src/cyhy_db/models/report_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy report documents.""" + # Standard Python Libraries from datetime import datetime from typing import List @@ -13,6 +15,8 @@ class ReportDoc(Document): + """The report document model.""" + model_config = ConfigDict(extra="forbid") owner: str @@ -21,6 +25,8 @@ class ReportDoc(Document): report_types: List[ReportType] class Settings: + """Beanie settings.""" + name = "reports" indexes = [ IndexModel( diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index cec041a..b8b259e 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy request documents.""" + # Standard Python Libraries from datetime import datetime from ipaddress import IPv4Network @@ -24,6 +26,8 @@ class Contact(BaseModel): + """A point of contact for the entity.""" + model_config = ConfigDict(extra="forbid") email: EmailStr @@ -33,6 +37,8 @@ class Contact(BaseModel): class Location(BaseModel): + """A location with various geographical identifiers.""" + model_config = ConfigDict(extra="forbid") country_name: str @@ -47,6 +53,8 @@ class Location(BaseModel): class Agency(BaseModel): + """Model representing a CyHy-enrolled entity.""" + model_config = ConfigDict(extra="forbid") name: str @@ -57,6 +65,8 @@ class Agency(BaseModel): class ScanLimit(BaseModel): + """Scan limits for a specific scan type.""" + model_config = ConfigDict(extra="forbid") scan_type: ScanType = Field(..., alias="scanType") @@ -64,6 +74,8 @@ class ScanLimit(BaseModel): class Window(BaseModel): + """A day and time window for scheduling scans.""" + model_config = ConfigDict(extra="forbid") day: DayOfWeek = Field(default=DayOfWeek.SUNDAY) @@ -72,13 +84,15 @@ class Window(BaseModel): @field_validator("start") def validate_start(cls, v): - # Validate that the start time is in the format HH:MM:SS + """Validate that the start time is in the correct format.""" if not re.match(r"^\d{2}:\d{2}:\d{2}$", v): raise ValueError("Start time must be in the format HH:MM:SS") return v class RequestDoc(Document): + """The request document model.""" + model_config = ConfigDict(extra="forbid") id: str = Field(default=BOGUS_ID) # type: ignore[assignment] @@ -100,10 +114,11 @@ class RequestDoc(Document): @before_event(Insert, Replace, ValidateOnSave) async def set_id_to_acronym(self): - # Set the id to the agency acronym if it is the default value + """Set the id to the agency acronym if it is the default value.""" if self.id == BOGUS_ID: self.id = self.agency.acronym class Settings: - # Beanie settings + """Beanie settings.""" + name = "requests" diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 794401c..89173fc 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -18,6 +18,8 @@ class ScanDoc(Document): # TODO: Make this a BaseModel + """The scan document model.""" + # Validate on assignment so ip_int is recalculated as ip is set model_config = ConfigDict(extra="forbid", validate_assignment=True) @@ -31,12 +33,14 @@ class ScanDoc(Document): # TODO: Make this a BaseModel @model_validator(mode="before") def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: + """Calculate the integer representation of an IP address.""" # ip may still be string if it was just set values["ip_int"] = int(ip_address(values["ip"])) return values class Settings: - # Beanie settings + """Beanie settings.""" + name = "scandocs" indexes = [ IndexModel( @@ -49,6 +53,7 @@ class Settings: @classmethod async def reset_latest_flag_by_owner(cls, owner: str): + """Reset the latest flag for all scans for a given owner.""" # flake8 E712 is "comparison to True should be 'if cond is True:' or 'if # cond:'" but this is unavoidable due to Beanie syntax. await cls.find(cls.latest == True, cls.owner == owner).update_many( # noqa E712 @@ -67,6 +72,7 @@ async def reset_latest_flag_by_ip( | str ), ): + """Reset the latest flag for all scans for a given IP address.""" if isinstance(ips, Iterable): # TODO Figure out why coverage thinks this next line can exit early ip_ints = [int(ip_address(x)) for x in ips] @@ -83,6 +89,7 @@ async def reset_latest_flag_by_ip( async def tag_latest( cls, owners: List[str], snapshot: Union["SnapshotDoc", ObjectId, str] ): + """Tag the latest scan for given owners with a snapshot id.""" from . import SnapshotDoc if isinstance(snapshot, SnapshotDoc): diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index 4f76d0b..dddc431 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy snapshot documents.""" + # Standard Python Libraries from datetime import datetime from ipaddress import IPv4Network @@ -12,6 +14,8 @@ class VulnerabilityCounts(BaseModel): + """The model for vulnerability counts.""" + model_config = ConfigDict(extra="forbid") critical: int = 0 @@ -22,6 +26,8 @@ class VulnerabilityCounts(BaseModel): class WorldData(BaseModel): + """The model for aggregated metrics of all CyHy entities.""" + model_config = ConfigDict(extra="forbid") host_count: int = 0 @@ -35,6 +41,8 @@ class WorldData(BaseModel): class TicketMetrics(BaseModel): + """The model for ticket metrics.""" + model_config = ConfigDict(extra="forbid") median: int = 0 @@ -42,6 +50,8 @@ class TicketMetrics(BaseModel): class TicketOpenMetrics(BaseModel): + """The model for open ticket metrics.""" + model_config = ConfigDict(extra="forbid") # Numbers in this section refer to how long open tix were open AT this date/time @@ -53,6 +63,8 @@ class TicketOpenMetrics(BaseModel): class TicketCloseMetrics(BaseModel): + """The model for closed ticket metrics.""" + model_config = ConfigDict(extra="forbid") # Numbers in this section only include tix that closed AT/AFTER this date/time @@ -64,6 +76,8 @@ class TicketCloseMetrics(BaseModel): class SnapshotDoc(Document): + """The snapshot document model.""" + model_config = ConfigDict(extra="forbid") owner: str = Field(...) @@ -91,7 +105,8 @@ class SnapshotDoc(Document): tix_msec_to_close: TicketCloseMetrics = Field(default_factory=TicketCloseMetrics) class Settings: - # Beanie settings + """Beanie settings.""" + name = "snapshots" indexes = [ IndexModel( diff --git a/src/cyhy_db/models/system_control_doc.py b/src/cyhy_db/models/system_control_doc.py index 606f575..fd0714e 100644 --- a/src/cyhy_db/models/system_control_doc.py +++ b/src/cyhy_db/models/system_control_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy system control documents.""" + # Standard Python Libraries import asyncio from datetime import datetime @@ -14,6 +16,8 @@ class SystemControlDoc(Document): + """The system control document model.""" + model_config = ConfigDict(extra="forbid") action: ControlAction @@ -24,6 +28,8 @@ class SystemControlDoc(Document): completed: bool = False # Set to True when after the action has occurred class Settings: + """Beanie settings.""" + name = "control" @classmethod diff --git a/src/cyhy_db/models/tally_doc.py b/src/cyhy_db/models/tally_doc.py index 37c1a57..d1a69af 100644 --- a/src/cyhy_db/models/tally_doc.py +++ b/src/cyhy_db/models/tally_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy tally documents.""" + # Standard Python Libraries from datetime import datetime @@ -9,6 +11,8 @@ class StatusCounts(BaseModel): + """The model for host status counts.""" + model_config = ConfigDict(extra="forbid") DONE: int = 0 @@ -18,6 +22,8 @@ class StatusCounts(BaseModel): class Counts(BaseModel): + """The model for scan stage counts.""" + model_config = ConfigDict(extra="forbid") BASESCAN: StatusCounts = Field(default_factory=StatusCounts) @@ -28,6 +34,8 @@ class Counts(BaseModel): class TallyDoc(Document): + """The tally document model.""" + model_config = ConfigDict(extra="forbid") _id: str # owner_id @@ -36,7 +44,10 @@ class TallyDoc(Document): @before_event(Insert, Replace, ValidateOnSave) async def before_save(self): + """Set data just prior to saving a tally document.""" self.last_change = utcnow() class Settings: + """Beanie settings.""" + name = "tallies" diff --git a/src/cyhy_db/models/vuln_scan_doc.py b/src/cyhy_db/models/vuln_scan_doc.py index 0b6ac6c..9060a5e 100644 --- a/src/cyhy_db/models/vuln_scan_doc.py +++ b/src/cyhy_db/models/vuln_scan_doc.py @@ -1,3 +1,5 @@ +"""The model for CyHy vulnerability scan documents.""" + # Standard Python Libraries from datetime import datetime @@ -10,6 +12,8 @@ class VulnScanDoc(ScanDoc): + """The vulnerability scan document model.""" + model_config = ConfigDict(extra="forbid") protocol: Protocol @@ -31,7 +35,8 @@ class VulnScanDoc(ScanDoc): synopsis: str class Settings: - # Beanie settings + """Beanie settings.""" + name = "vuln_scans" indexes = [ IndexModel( diff --git a/src/cyhy_db/utils/decorators.py b/src/cyhy_db/utils/decorators.py index 83d20c1..3c64523 100644 --- a/src/cyhy_db/utils/decorators.py +++ b/src/cyhy_db/utils/decorators.py @@ -1,8 +1,12 @@ +"""Decorators for the cyhy_db package.""" + # Standard Python Libraries import warnings def deprecated(reason): + """Mark a function as deprecated.""" + def decorator(func): if isinstance(reason, str): message = f"{func.__name__} is deprecated and will be removed in a future version. {reason}" diff --git a/src/cyhy_db/utils/time.py b/src/cyhy_db/utils/time.py index 683e660..3997df2 100644 --- a/src/cyhy_db/utils/time.py +++ b/src/cyhy_db/utils/time.py @@ -1,3 +1,5 @@ +"""Utility functions for working with time and dates.""" + # Standard Python Libraries from datetime import datetime, timezone diff --git a/tests/conftest.py b/tests/conftest.py index be2c53e..418f235 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -92,6 +92,7 @@ def mongodb_container(docker_client, mongo_image_tag): @pytest.fixture(autouse=True, scope="session") def mongo_express_container(docker_client, request): + """Fixture for the Mongo Express test container.""" if not request.config.getoption("--mongo-express"): yield None return diff --git a/tests/test_connection.py b/tests/test_connection.py index 0fe7564..aae24f2 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -9,6 +9,7 @@ async def test_connection_motor(db_uri, db_name): + """Test connectivity to database.""" client = AsyncIOMotorClient(db_uri) db = client[db_name] server_info = await db.command("ping") @@ -16,6 +17,7 @@ async def test_connection_motor(db_uri, db_name): async def test_connection_beanie(): + """Test a simple database query.""" # Attempt to find a document in the empty CVE collection # await initialize_db(db_uri, db_name) # Manually initialize for testing result = await CVE.get("CVE-2024-DOES-NOT-EXIST") diff --git a/tests/test_data_generator.py b/tests/test_data_generator.py index 2ee0bb5..81c79a0 100644 --- a/tests/test_data_generator.py +++ b/tests/test_data_generator.py @@ -1,3 +1,11 @@ +""" +This module generates test data for CyHy reports using factory classes. + +It includes factories for creating instances of various models such as CVE, +Agency, Contact, Location, Window, and RequestDoc. Additionally, it provides a +custom provider for generating specific data like CVE IDs and IPv4 networks. +""" + # Standard Python Libraries from datetime import datetime import ipaddress @@ -27,25 +35,38 @@ class CyHyProvider(BaseProvider): + """Custom provider for generating specific CyHy data.""" + class Meta: + """Meta class for CyHyProvider.""" + name = "cyhy_provider" def cve_id(self, year=None): - # If year is None, generate a random year between 1999 and the current year + """ + Generate a CVE ID. + + Args: + year (int, optional): The year for the CVE ID. If None, a random + year between 1999 and the current year is used. + + Returns: + str: A CVE ID in the format CVE-YYYY-NNNNN. + """ if year is None: year = self.random.randint(1999, datetime.now().year) - - # Generate a random number for the CVE, ensuring it has a leading zero if necessary number = self.random.randint(1, 99999) - return f"CVE-{year}-{number:05d}" def network_ipv4(self): - # Generate a base IP address + """ + Generate an IPv4 network. + + Returns: + ipaddress.IPv4Network: A randomly generated IPv4 network. + """ base_ip = generic.internet.ip_v4() - # Choose a random CIDR between 24-30 to ensure a smaller network size and avoid host bits set error cidr = random.randint(24, 30) - # Create the network address network = ipaddress.IPv4Network(f"{base_ip}/{cidr}", strict=False) return network @@ -56,7 +77,11 @@ def network_ipv4(self): @register class CVEFactory(factory.Factory): + """Factory for creating CVE instances.""" + class Meta: + """Meta class for CVEFactory.""" + model = CVE id = factory.LazyFunction(lambda: generic.cyhy_provider.cve_id()) @@ -66,11 +91,14 @@ class Meta: class AgencyFactory(factory.Factory): + """Factory for creating Agency instances.""" + class Meta: + """Meta class for AgencyFactory.""" + model = Agency name = factory.Faker("company") - # Generate an acronym from the name acronym = factory.LazyAttribute( lambda o: "".join(word[0].upper() for word in o.name.split()) ) @@ -82,7 +110,11 @@ class Meta: class ContactFactory(factory.Factory): + """Factory for creating Contact instances.""" + class Meta: + """Meta class for ContactFactory.""" + model = Contact email = factory.Faker("email") @@ -92,7 +124,11 @@ class Meta: class LocationFactory(factory.Factory): + """Factory for creating Location instances.""" + class Meta: + """Meta class for LocationFactory.""" + model = Location country_name = factory.Faker("country") @@ -107,7 +143,11 @@ class Meta: class WindowFactory(factory.Factory): + """Factory for creating Window instances.""" + class Meta: + """Meta class for WindowFactory.""" + model = Window day = factory.LazyFunction(lambda: random.choice(list(DayOfWeek))) @@ -116,7 +156,11 @@ class Meta: class RequestDocFactory(factory.Factory): + """Factory for creating RequestDoc instances.""" + class Meta: + """Meta class for RequestDocFactory.""" + model = RequestDoc id = factory.LazyAttribute( @@ -139,13 +183,13 @@ class Meta: generic.cyhy_provider.network_ipv4() for _ in range(random.randint(1, 5)) ] ) - # create a set of 1 to 3 random scan types from the ScanType enum scan_types = factory.LazyFunction( lambda: {random.choice(list(ScanType)) for _ in range(random.randint(1, 3))} ) async def test_create_cves(): + """Test function to create and save 100 CVE instances.""" for _ in range(100): cve = CVEFactory() print(cve) @@ -153,6 +197,7 @@ async def test_create_cves(): async def test_create_request_docs(): + """Test function to create and save 100 RequestDoc instances.""" for _ in range(100): request_doc = RequestDocFactory() print(request_doc) diff --git a/tests/test_host_doc.py b/tests/test_host_doc.py index 73130bf..6fb1368 100644 --- a/tests/test_host_doc.py +++ b/tests/test_host_doc.py @@ -1,3 +1,5 @@ +"""Test HostDoc model functionality.""" + # Standard Python Libraries from ipaddress import ip_address @@ -11,6 +13,7 @@ def test_host_doc_init(): + """Test HostDoc object initialization.""" # Create a HostDoc object host_doc = HostDoc( ip=ip_address(VALID_IP_1_STR), @@ -21,6 +24,7 @@ def test_host_doc_init(): async def test_save(): + """Test saving a HostDoc object to the database.""" # Create a HostDoc object host_doc = HostDoc( ip=ip_address(VALID_IP_1_STR), @@ -32,6 +36,7 @@ async def test_save(): async def test_get_by_ip(): + """Test finding a HostDoc object by its IP address.""" # Find a HostDoc object by its IP address host_doc = await HostDoc.get_by_ip(ip_address(VALID_IP_1_STR)) assert host_doc.ip == ip_address(VALID_IP_1_STR) diff --git a/tests/test_request_doc.py b/tests/test_request_doc.py index 95d511d..4adbb3a 100644 --- a/tests/test_request_doc.py +++ b/tests/test_request_doc.py @@ -6,8 +6,8 @@ async def test_init(): + """Test RequestDoc object initialization.""" # Create a RequestDoc object - request_doc = RequestDoc( agency=Agency( name="Cybersecurity and Infrastructure Security Agency", acronym="CISA" diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index 682c375..def4c1e 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -18,6 +18,11 @@ def test_ip_int_init(): + """Test IP address integer conversion on initialization. + + This test verifies that the IP address is correctly converted to an + integer when a ScanDoc object is initialized. + """ # Create a ScanDoc object scan_doc = ScanDoc( ip=ipaddress.ip_address(VALID_IP_1_STR), @@ -31,6 +36,11 @@ def test_ip_int_init(): def test_ip_int_change(): + """Test IP address integer conversion on IP address change. + + This test verifies that the IP address is correctly converted to an + integer when the IP address of a ScanDoc object is changed. + """ # Create a ScanDoc object scan_doc = ScanDoc( ip=ipaddress.ip_address(VALID_IP_1_STR), @@ -46,6 +56,12 @@ def test_ip_int_change(): def test_ip_string_set(): + """Test IP address string conversion and integer calculation. + + This test verifies that an IP address provided as a string is correctly + converted to an ipaddress.IPv4Address object and that the corresponding + integer value is calculated correctly. + """ scan_doc = ScanDoc( ip=VALID_IP_1_STR, owner="YOUR_MOM", @@ -59,6 +75,11 @@ def test_ip_string_set(): async def test_ip_address_field_fetch(): + """Test IP address retrieval from the database. + + This test verifies that the IP address of a ScanDoc object is correctly + retrieved from the database. + """ # Create a ScanDoc object scan_doc = ScanDoc( ip=ipaddress.ip_address(VALID_IP_1_STR), @@ -81,6 +102,11 @@ async def test_ip_address_field_fetch(): def test_invalid_ip_address(): + """Test validation error for invalid IP addresses. + + This test verifies that a ValidationError is raised when an invalid IP + address is provided to a ScanDoc object. + """ with pytest.raises(ValidationError): ScanDoc( ip="999.999.999.999", # This should be invalid @@ -90,6 +116,11 @@ def test_invalid_ip_address(): async def test_reset_latest_flag_by_owner(): + """Test resetting the latest flag by owner. + + This test verifies that the latest flag of ScanDoc objects is correctly + reset when the reset_latest_flag_by_owner method is called. + """ # Create a ScanDoc object OWNER = "RESET_BY_OWNER" scan_doc = ScanDoc( @@ -107,7 +138,11 @@ async def test_reset_latest_flag_by_owner(): async def test_reset_latest_flag_by_ip(): - # Create a ScanDoc object + """Test resetting the latest flag by IP address. + + This test verifies that the latest flag of ScanDoc objects is correctly + reset when the reset_latest_flag_by_ip method is called. + """ IP_TO_RESET_1 = ipaddress.ip_address("128.205.1.2") IP_TO_RESET_2 = ipaddress.ip_address("128.205.1.3") scan_doc_1 = ScanDoc(ip=IP_TO_RESET_1, owner="RESET_BY_IP", source="nmap") @@ -131,8 +166,12 @@ async def test_reset_latest_flag_by_ip(): async def test_tag_latest(): - # Create a SnapshotDoc object + """Test tagging the latest scan with a snapshot. + This test verifies that the latest ScanDoc object is correctly tagged + with a SnapshotDoc object when the tag_latest method is called. + """ + # Create a SnapshotDoc object owner = "TAG_LATEST" snapshot_doc = SnapshotDoc( owner=owner, From 7b6e816c21dbf1d01cd0be17dbf1af515e7e1b6c Mon Sep 17 00:00:00 2001 From: David Redmin Date: Fri, 30 Aug 2024 16:04:52 -0400 Subject: [PATCH 079/139] Tell bandit to ignore warnings about our use of random numbers We aren't using Random() for the purposes of cryptography here, so we can safely ignore these warnings. --- tests/test_data_generator.py | 74 +++++++++++++++++++++++++++--------- 1 file changed, 56 insertions(+), 18 deletions(-) diff --git a/tests/test_data_generator.py b/tests/test_data_generator.py index 81c79a0..9edffe9 100644 --- a/tests/test_data_generator.py +++ b/tests/test_data_generator.py @@ -66,7 +66,11 @@ def network_ipv4(self): ipaddress.IPv4Network: A randomly generated IPv4 network. """ base_ip = generic.internet.ip_v4() - cidr = random.randint(24, 30) + # The following line generates a warning from bandit about "Standard + # pseudo-random generators are not suitable for security/cryptographic + # purposes." We aren't using Random() for the purposes of cryptography + # here, so we can safely ignore that warning. + cidr = random.randint(24, 30) # nosec B311 network = ipaddress.IPv4Network(f"{base_ip}/{cidr}", strict=False) return network @@ -85,9 +89,17 @@ class Meta: model = CVE id = factory.LazyFunction(lambda: generic.cyhy_provider.cve_id()) - cvss_score = factory.LazyFunction(lambda: round(random.uniform(0, 10), 1)) - cvss_version = factory.LazyFunction(lambda: random.choice(list(CVSSVersion))) - severity = factory.LazyFunction(lambda: random.randint(1, 4)) + # The following lines generate warnings from bandit about "Standard + # pseudo-random generators are not suitable for security/cryptographic + # purposes." We aren't using Random() for the purposes of cryptography + # here, so we can safely ignore those warnings. + cvss_score = factory.LazyFunction( + lambda: round(random.uniform(0, 10), 1) # nosec B311 + ) + cvss_version = factory.LazyFunction( + lambda: random.choice(list(CVSSVersion)) # nosec B311 + ) + severity = factory.LazyFunction(lambda: random.randint(1, 4)) # nosec B311 class AgencyFactory(factory.Factory): @@ -102,9 +114,13 @@ class Meta: acronym = factory.LazyAttribute( lambda o: "".join(word[0].upper() for word in o.name.split()) ) - type = factory.LazyFunction(lambda: random.choice(list(AgencyType))) + # The following lines generate warnings from bandit about "Standard + # pseudo-random generators are not suitable for security/cryptographic + # purposes." We aren't using Random() for the purposes of cryptography + # here, so we can safely ignore those warnings. + type = factory.LazyFunction(lambda: random.choice(list(AgencyType))) # nosec B311 contacts = factory.LazyFunction( - lambda: [ContactFactory() for _ in range(random.randint(1, 5))] + lambda: [ContactFactory() for _ in range(random.randint(1, 5))] # nosec B311 ) location = factory.LazyFunction(lambda: LocationFactory()) @@ -120,7 +136,11 @@ class Meta: email = factory.Faker("email") name = factory.Faker("name") phone = factory.Faker("phone_number") - type = factory.LazyFunction(lambda: random.choice(list(PocType))) + # The following line generates a warning from bandit about "Standard + # pseudo-random generators are not suitable for security/cryptographic + # purposes." We aren't using Random() for the purposes of cryptography + # here, so we can safely ignore that warning. + type = factory.LazyFunction(lambda: random.choice(list(PocType))) # nosec B311 class LocationFactory(factory.Factory): @@ -150,8 +170,12 @@ class Meta: model = Window - day = factory.LazyFunction(lambda: random.choice(list(DayOfWeek))) - duration = factory.LazyFunction(lambda: random.randint(0, 168)) + # The following lines generate warnings from bandit about "Standard + # pseudo-random generators are not suitable for security/cryptographic + # purposes." We aren't using Random() for the purposes of cryptography + # here, so we can safely ignore those warnings. + day = factory.LazyFunction(lambda: random.choice(list(DayOfWeek))) # nosec B311 + duration = factory.LazyFunction(lambda: random.randint(0, 168)) # nosec B311 start = factory.Faker("time", pattern="%H:%M:%S") @@ -163,28 +187,42 @@ class Meta: model = RequestDoc + # The following lines generate warnings from bandit about "Standard + # pseudo-random generators are not suitable for security/cryptographic + # purposes." We aren't using Random() for the purposes of cryptography + # here, so we can safely ignore those warnings. id = factory.LazyAttribute( - lambda o: o.agency.acronym + "-" + str(random.randint(1, 1000)) + lambda o: o.agency.acronym + "-" + str(random.randint(1, 1000)) # nosec B311 ) agency = factory.SubFactory(AgencyFactory) enrolled = factory.LazyFunction(utcnow) - init_stage = factory.LazyFunction(lambda: random.choice(list(Stage))) + init_stage = factory.LazyFunction(lambda: random.choice(list(Stage))) # nosec B311 key = factory.Faker("password") period_start = factory.LazyFunction(utcnow) - report_period = factory.LazyFunction(lambda: random.choice(list(ReportPeriod))) - retired = factory.LazyFunction(lambda: random.choice([True, False])) - scheduler = factory.LazyFunction(lambda: random.choice(list(Scheduler))) - stakeholder = factory.LazyFunction(lambda: random.choice([True, False])) + report_period = factory.LazyFunction( + lambda: random.choice(list(ReportPeriod)) # nosec B311 + ) + retired = factory.LazyFunction(lambda: random.choice([True, False])) # nosec B311 + scheduler = factory.LazyFunction( + lambda: random.choice(list(Scheduler)) # nosec B311 + ) + stakeholder = factory.LazyFunction( + lambda: random.choice([True, False]) # nosec B311 + ) windows = factory.LazyFunction( - lambda: [WindowFactory() for _ in range(random.randint(1, 5))] + lambda: [WindowFactory() for _ in range(random.randint(1, 5))] # nosec B311 ) networks = factory.LazyFunction( lambda: [ - generic.cyhy_provider.network_ipv4() for _ in range(random.randint(1, 5)) + generic.cyhy_provider.network_ipv4() + for _ in range(random.randint(1, 5)) # nosec B311 ] ) scan_types = factory.LazyFunction( - lambda: {random.choice(list(ScanType)) for _ in range(random.randint(1, 3))} + lambda: { + random.choice(list(ScanType)) # nosec B311 + for _ in range(random.randint(1, 3)) # nosec B311 + } ) From 9b828ac7b432eb5ba7558b4540d21d5718aeba93 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 25 Sep 2024 17:50:47 -0400 Subject: [PATCH 080/139] Limit Python version to 3.10 and greater We are using the or (`|`) operator for types that was introduced in Python 3.10 with PEP 604. See: - https://docs.python.org/3/whatsnew/3.10.html#pep-604-new-type-union-operator --- .github/workflows/build.yml | 6 ------ setup.py | 4 +--- 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 3a7ef12..5cc0632 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -180,8 +180,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.8" - - "3.9" - "3.10" - "3.11" - "3.12" @@ -285,8 +283,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.8" - - "3.9" - "3.10" - "3.11" - "3.12" @@ -339,8 +335,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.8" - - "3.9" - "3.10" - "3.11" - "3.12" diff --git a/setup.py b/setup.py index edde11e..7c01781 100644 --- a/setup.py +++ b/setup.py @@ -75,14 +75,12 @@ def get_version(version_file): # that you indicate whether you support Python 2, Python 3 or both. "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: Implementation :: CPython", ], - python_requires=">=3.8", + python_requires=">=3.10", # What does your project relate to? keywords=["cyhy", "database"], packages=find_packages(where="src"), From fa658d012345b0878a5a6fe33d3a5c00976704a1 Mon Sep 17 00:00:00 2001 From: Felddy Date: Wed, 25 Sep 2024 17:52:37 -0400 Subject: [PATCH 081/139] Uncomment dependabot ignores handled by skeleton --- .github/dependabot.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 9e4ff7b..e3d24b7 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -19,10 +19,10 @@ updates: - dependency-name: hashicorp/setup-terraform - dependency-name: mxschmitt/action-tmate - dependency-name: step-security/harden-runner - # # Managed by cisagov/cyhy-db - # - dependency-name: actions/download-artifact - # - dependency-name: actions/upload-artifact - # - dependency-name: github/codeql-action + # Managed by cisagov/skeleton-python-library + - dependency-name: actions/download-artifact + - dependency-name: actions/upload-artifact + - dependency-name: github/codeql-action package-ecosystem: github-actions schedule: interval: weekly From 030607c858847efd9c00fe540c5bf11b9a917787 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 1 Oct 2024 16:33:07 -0400 Subject: [PATCH 082/139] Explicitly set the pytest log level to INFO This is primarily because pymongo DEBUG level logging is extremely verbose and should only be enabled when needed. Co-authored-by: Mark Feldhousen --- pytest.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytest.ini b/pytest.ini index 189f9fc..ae07d53 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,3 +1,3 @@ [pytest] -addopts = -v -ra --cov +addopts = -v -ra --cov --log-cli-level=INFO asyncio_mode = auto From 9aad89664d3b4dd1bba5245f9e8ade53cdc37aec Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 1 Oct 2024 16:37:42 -0400 Subject: [PATCH 083/139] Restore a useful comment and import from the skeleton Co-authored-by: Mark Feldhousen --- src/cyhy_db/__init__.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/cyhy_db/__init__.py b/src/cyhy_db/__init__.py index 4b9e59d..36ea325 100644 --- a/src/cyhy_db/__init__.py +++ b/src/cyhy_db/__init__.py @@ -1,5 +1,11 @@ """The cyhy_db package provides an interface to a CyHy database.""" +# We disable a Flake8 check for "Module imported but unused (F401)" here because +# although this import is not directly used, it populates the value +# package_name.__version__, which is used to get version information about this +# Python package. + +from ._version import __version__ # noqa: F401 from .db import initialize_db __all__ = ["initialize_db"] From 86d8c38fd4aa3d5b03febc4dd9b8ba9a0f254130 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 1 Oct 2024 16:38:30 -0400 Subject: [PATCH 084/139] Remove a commented out line that is no longer needed Co-authored-by: Mark Feldhousen --- tests/test_connection.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_connection.py b/tests/test_connection.py index aae24f2..27ca0b0 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -19,6 +19,5 @@ async def test_connection_motor(db_uri, db_name): async def test_connection_beanie(): """Test a simple database query.""" # Attempt to find a document in the empty CVE collection - # await initialize_db(db_uri, db_name) # Manually initialize for testing result = await CVE.get("CVE-2024-DOES-NOT-EXIST") assert result is None, "Expected no document to be found" From 0afb2ed78179a838cdf29666785d0dc97f3633e7 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 1 Oct 2024 16:41:32 -0400 Subject: [PATCH 085/139] Add two unit tests for the initialize_db function Co-authored-by: Mark Feldhousen --- tests/test_connection.py | 45 ++++++++++++++++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/tests/test_connection.py b/tests/test_connection.py index 27ca0b0..4febd58 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -1,10 +1,14 @@ """Test database connection.""" +# Standard Python Libraries +from unittest.mock import AsyncMock, patch + # Third-Party Libraries -from motor.motor_asyncio import AsyncIOMotorClient +from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase +import pytest # cisagov Libraries -# from cyhy_db import initialize_db +from cyhy_db.db import ALL_MODELS, initialize_db from cyhy_db.models import CVE @@ -21,3 +25,40 @@ async def test_connection_beanie(): # Attempt to find a document in the empty CVE collection result = await CVE.get("CVE-2024-DOES-NOT-EXIST") assert result is None, "Expected no document to be found" + + +# Confused about the order of patch statements relative to the order of the test +# function parameters? See here: +# https://docs.python.org/3/library/unittest.mock.html#patch +@patch("cyhy_db.db.init_beanie", return_value=None) # mock_init_beanie +@patch("cyhy_db.db.AsyncIOMotorClient") # mock_async_iomotor_client +@pytest.mark.asyncio +async def test_initialize_db_success(mock_async_iomotor_client, mock_init_beanie): + """Test a success case of the initialize_db function.""" + db_uri = "mongodb://localhost:27017" + db_name = "test_db" + + mock_client = AsyncMock() + mock_db = AsyncMock(AsyncIOMotorDatabase) + mock_client.__getitem__.return_value = mock_db + mock_async_iomotor_client.return_value = mock_client + + db = await initialize_db(db_uri, db_name) + assert db == mock_db + mock_async_iomotor_client.assert_called_once_with(db_uri) + mock_init_beanie.assert_called_once_with( + database=mock_db, document_models=ALL_MODELS + ) + + +@pytest.mark.asyncio +async def test_initialize_db_failure(): + """Test a failure case of the initialize_db function.""" + db_uri = "mongodb://localhost:27017" + db_name = "test_db" + + with patch( + "cyhy_db.db.AsyncIOMotorClient", side_effect=Exception("Connection error") + ): + with pytest.raises(Exception, match="Connection error"): + await initialize_db(db_uri, db_name) From ddab65e7fc98aaad0527f111a5ffa866c9d1d538 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 1 Oct 2024 22:46:04 -0400 Subject: [PATCH 086/139] Correctly set state in HostDoc --- src/cyhy_db/models/host_doc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 1f28eeb..6989d55 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -114,11 +114,11 @@ def set_state(self, nmap_says_up, has_open_ports, reason=None): either argument can be None. """ if has_open_ports: # Only PORTSCAN sends in has_open_ports - self.state = State(True, "open-port") + self.state = State(up=True, reason="open-port") elif has_open_ports is False: - self.state = State(False, "no-open") + self.state = State(up=False, reason="no-open") elif nmap_says_up is False: # NETSCAN says host is down - self.state = State(False, reason) + self.state = State(up=False, reason=reason) # TODO: There are a lot of functions in the Python 2 version that may or may not be used. # Instead of porting them all over, we should just port them as they are needed. From 998315047467bc7638967bafd63918b3357189e7 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 1 Oct 2024 22:47:47 -0400 Subject: [PATCH 087/139] Add unit tests for HostDoc.set_state() --- tests/test_host_doc.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/test_host_doc.py b/tests/test_host_doc.py index 6fb1368..2201451 100644 --- a/tests/test_host_doc.py +++ b/tests/test_host_doc.py @@ -5,6 +5,7 @@ # cisagov Libraries from cyhy_db.models import HostDoc +from cyhy_db.models.host_doc import State VALID_IP_1_STR = "0.0.0.1" VALID_IP_2_STR = "0.0.0.2" @@ -40,3 +41,27 @@ async def test_get_by_ip(): # Find a HostDoc object by its IP address host_doc = await HostDoc.get_by_ip(ip_address(VALID_IP_1_STR)) assert host_doc.ip == ip_address(VALID_IP_1_STR) + + +async def test_set_state_open_ports(): + """Test setting HostDoc state with open ports.""" + # Find a HostDoc object by its IP address + host_doc = await HostDoc.get_by_ip(ip_address(VALID_IP_1_STR)) + host_doc.set_state(nmap_says_up=None, has_open_ports=True) + assert host_doc.state == State(up=True, reason="open-port") + + +async def test_set_state_no_open_ports(): + """Test setting HostDoc state with no open ports.""" + # Find a HostDoc object by its IP address + host_doc = await HostDoc.get_by_ip(ip_address(VALID_IP_1_STR)) + host_doc.set_state(nmap_says_up=None, has_open_ports=False) + assert host_doc.state == State(up=False, reason="no-open") + + +async def test_set_state_nmap_says_down(): + """Test setting HostDoc state when nmap says the host is down.""" + # Find a HostDoc object by its IP address + host_doc = await HostDoc.get_by_ip(ip_address(VALID_IP_1_STR)) + host_doc.set_state(nmap_says_up=False, has_open_ports=None, reason="no-reply") + assert host_doc.state == State(up=False, reason="no-reply") From b510c09a9a41062bd2ca0bea6d8463b15b6edcc0 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 1 Oct 2024 23:04:27 -0400 Subject: [PATCH 088/139] Add unit tests for SystemControlDoc --- tests/test_system_control_doc.py | 54 ++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 tests/test_system_control_doc.py diff --git a/tests/test_system_control_doc.py b/tests/test_system_control_doc.py new file mode 100644 index 0000000..a3cd1dc --- /dev/null +++ b/tests/test_system_control_doc.py @@ -0,0 +1,54 @@ +"""Test SystemControlDoc model functionality.""" + +# Standard Python Libraries +from datetime import timedelta +from unittest.mock import AsyncMock, patch + +# cisagov Libraries +from cyhy_db.models.system_control_doc import SystemControlDoc +from cyhy_db.utils import utcnow + + +async def test_wait_for_completion_completed(): + """Test wait_for_completion when the document is completed.""" + document_id = "test_id" + mock_doc = AsyncMock() + mock_doc.completed = True + + with patch.object(SystemControlDoc, "get", return_value=mock_doc): + result = await SystemControlDoc.wait_for_completion(document_id) + assert result is True + + +async def test_wait_for_completion_timeout(): + """Test wait_for_completion when the document is not completed before the timeout.""" + document_id = "test_id" + mock_doc = AsyncMock() + mock_doc.completed = False + + with patch.object(SystemControlDoc, "get", return_value=mock_doc): + with patch( + "cyhy_db.models.system_control_doc.utcnow", + side_effect=[utcnow(), utcnow() + timedelta(seconds=10)], + ): + result = await SystemControlDoc.wait_for_completion(document_id, timeout=5) + assert result is False + + +async def test_wait_for_completion_no_timeout(): + """Test wait_for_completion when a timeout is not set.""" + document_id = "test_id" + mock_doc = AsyncMock() + mock_doc.completed = False + + async def side_effect(*args, **kwargs): + if side_effect.call_count == 2: + mock_doc.completed = True + side_effect.call_count += 1 + return mock_doc + + side_effect.call_count = 0 + + with patch.object(SystemControlDoc, "get", side_effect=side_effect): + result = await SystemControlDoc.wait_for_completion(document_id) + assert result is True From 5642a7c10cf8718843878f0c50e0811e8d537055 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 1 Oct 2024 23:32:34 -0400 Subject: [PATCH 089/139] Add additional unit tests for ScanDoc.tag_latest() --- tests/test_scan_doc.py | 97 +++++++++++++++++++++++++++++++++++++++--- 1 file changed, 92 insertions(+), 5 deletions(-) diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index def4c1e..c605134 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -165,14 +165,14 @@ async def test_reset_latest_flag_by_ip(): assert scan_doc_2.latest is False -async def test_tag_latest(): - """Test tagging the latest scan with a snapshot. +async def test_tag_latest_snapshot_doc(): + """Test tagging the latest scan with a SnapshotDoc. - This test verifies that the latest ScanDoc object is correctly tagged - with a SnapshotDoc object when the tag_latest method is called. + This test verifies that the latest ScanDoc object is correctly tagged with a + SnapshotDoc object when the tag_latest method is called with a SnapshotDoc. """ # Create a SnapshotDoc object - owner = "TAG_LATEST" + owner = "TAG_LATEST_SNAPSHOT_DOC" snapshot_doc = SnapshotDoc( owner=owner, start_time=utcnow(), @@ -195,3 +195,90 @@ async def test_tag_latest(): # Check that the scan now has a snapshot assert scan_doc.snapshots == [snapshot_doc], "Snapshot not added to scan" + + +async def test_tag_latest_snapshot_id(): + """Test tagging the latest scan with a snapshot ObjectId. + + This test verifies that the latest ScanDoc object is correctly tagged with a + SnapshotDoc object when the tag_latest method is called with a snapshot + ObjectId. + """ + # Create a SnapshotDoc object + owner = "TAG_LATEST_SNAPSHOT_ID" + snapshot_doc = SnapshotDoc( + owner=owner, + start_time=utcnow(), + end_time=utcnow(), + ) + await snapshot_doc.save() + # Create a ScanDoc object + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner=owner, + source="nmap", + ) + await scan_doc.save() + + # Tag the latest scan with the snapshot id + await ScanDoc.tag_latest([owner], snapshot_doc.id) + + # Retrieve the ScanDoc object from the database + scan_doc = await ScanDoc.find_one(ScanDoc.id == scan_doc.id, fetch_links=True) + + # Check that the scan now has a snapshot + assert scan_doc.snapshots == [snapshot_doc], "Snapshot not added to scan" + + +async def test_tag_latest_snapshot_id_str(): + """Test tagging the latest scan with the string representation of a snapshot ObjectId. + + This test verifies that the latest ScanDoc object is correctly tagged with a + SnapshotDoc object when the tag_latest method is called with the string + representation of a snapshot ObjectId. + """ + # Create a SnapshotDoc object + owner = "TAG_LATEST_SNAPSHOT_ID_STR" + snapshot_doc = SnapshotDoc( + owner=owner, + start_time=utcnow(), + end_time=utcnow(), + ) + await snapshot_doc.save() + # Create a ScanDoc object + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner=owner, + source="nmap", + ) + await scan_doc.save() + + # Tag the latest scan with the string representation of the snapshot id + await ScanDoc.tag_latest([owner], str(snapshot_doc.id)) + + # Retrieve the ScanDoc object from the database + scan_doc = await ScanDoc.find_one(ScanDoc.id == scan_doc.id, fetch_links=True) + + # Check that the scan now has a snapshot + assert scan_doc.snapshots == [snapshot_doc], "Snapshot not added to scan" + + +async def test_tag_latest_invalid_type(): + """Test tagging the latest scan with an invalid object type.""" + owner = "TAG_LATEST_INVALID_TYPE" + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner=owner, + source="nmap", + ) + await scan_doc.save() + + with pytest.raises(ValueError, match="Invalid snapshot type"): + # Attempt to tag the latest scan with an invalid object type + await ScanDoc.tag_latest([owner], 12345) + + # Retrieve the ScanDoc object from the database + scan_doc = await ScanDoc.find_one(ScanDoc.id == scan_doc.id, fetch_links=True) + + # Confirm that the scan does not have a snapshot + assert scan_doc.snapshots == [], "Scan should not have any snapshots" From 2870a9a066bb0a95f724b7f1916edc157ccb3ccf Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 10:02:26 -0400 Subject: [PATCH 090/139] Add a "no op" unit test for HostDoc.set_state() --- tests/test_host_doc.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_host_doc.py b/tests/test_host_doc.py index 2201451..ef20b54 100644 --- a/tests/test_host_doc.py +++ b/tests/test_host_doc.py @@ -65,3 +65,25 @@ async def test_set_state_nmap_says_down(): host_doc = await HostDoc.get_by_ip(ip_address(VALID_IP_1_STR)) host_doc.set_state(nmap_says_up=False, has_open_ports=None, reason="no-reply") assert host_doc.state == State(up=False, reason="no-reply") + + +async def test_set_state_no_op(): + """Test setting HostDoc state when inputs are supplied that results in no state change.""" + # Create a HostDoc object + host_doc = HostDoc( + ip=ip_address(VALID_IP_2_STR), + owner="NO-OP", + ) + # Save the HostDoc object to the database + await host_doc.save() + assert host_doc.id == VALID_IP_2_INT + + # Find HostDoc object by its IP address + host_doc = await HostDoc.get_by_ip(ip_address(VALID_IP_2_INT)) + assert host_doc.state == State(up=False, reason="new") + + host_doc.set_state(nmap_says_up=True, has_open_ports=None, reason="no-op-test-1") + assert host_doc.state == State(up=False, reason="new") + + host_doc.set_state(nmap_says_up=None, has_open_ports=None, reason="no-op-test-2") + assert host_doc.state == State(up=False, reason="new") From 41fefe8a5fb442daf975f31de5d645d87d4daa6f Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 10:24:33 -0400 Subject: [PATCH 091/139] Add unit tests for utils/decorators.py --- tests/test_decorators.py | 39 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 tests/test_decorators.py diff --git a/tests/test_decorators.py b/tests/test_decorators.py new file mode 100644 index 0000000..3f02638 --- /dev/null +++ b/tests/test_decorators.py @@ -0,0 +1,39 @@ +"""Test utils/decorators.py functionality.""" + +# Third-Party Libraries +import pytest + +# cisagov Libraries +from cyhy_db.utils.decorators import deprecated + + +def test_deprecated_decorator_with_reason(): + """Test the deprecated decorator with a reason.""" + + @deprecated("Use another function") + def old_function(): + """Impersonate a deprecated function.""" + return "result" + + with pytest.warns( + DeprecationWarning, + match="old_function is deprecated and will be removed in a future version. Use another function", + ): + result = old_function() + assert result == "result" + + +def test_deprecated_decorator_without_reason(): + """Test the deprecated decorator without a reason.""" + + @deprecated(None) + def old_function(): + """Impersonate a deprecated function.""" + return "result" + + with pytest.warns( + DeprecationWarning, + match="old_function is deprecated and will be removed in a future version.", + ): + result = old_function() + assert result == "result" From 323666eb450d5f319f65bfd2ee080896d283d933 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 14:39:07 -0400 Subject: [PATCH 092/139] Improve field validator for start window in RequestDoc This field now accepts either a string representing the time (%H:%M:%S) or a datetime.time object, which is stored in the database as a string in %H:%M:%S format. This also improves our input checking so that invalid times (e.g. 34:45:56) can no longer be added. Co-authored-by: Mark Feldhousen --- src/cyhy_db/models/request_doc.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index b8b259e..e729723 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -1,9 +1,8 @@ """The model for CyHy request documents.""" # Standard Python Libraries -from datetime import datetime +from datetime import datetime, time from ipaddress import IPv4Network -import re from typing import List, Optional # Third-Party Libraries @@ -80,14 +79,21 @@ class Window(BaseModel): day: DayOfWeek = Field(default=DayOfWeek.SUNDAY) duration: int = Field(default=168, ge=0, le=168) - start: str = Field(default="00:00:00") - - @field_validator("start") - def validate_start(cls, v): - """Validate that the start time is in the correct format.""" - if not re.match(r"^\d{2}:\d{2}:\d{2}$", v): - raise ValueError("Start time must be in the format HH:MM:SS") - return v + start: time = Field(default=time(0, 0, 0)) + + @field_validator("start", mode="before") + @classmethod + def parse_time(cls, v): + """Parse and validate a time representation.""" + if isinstance(v, str): + # Parse the string to datetime.time + return datetime.strptime(v, "%H:%M:%S").time() + elif isinstance(v, time): + return v + else: + raise ValueError( + "Invalid time format. Expected a string in '%H:%M:%S' format or datetime.time instance." + ) class RequestDoc(Document): @@ -121,4 +127,7 @@ async def set_id_to_acronym(self): class Settings: """Beanie settings.""" + bson_encoders = { + time: lambda value: value.strftime("%H:%M:%S") + } # Register custom encoder for datetime.time name = "requests" From 3b76e7bf78d8fff3a6b54f0b30ec15d3297107f8 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 14:41:22 -0400 Subject: [PATCH 093/139] Add unit tests for parse_time in Window (request_doc.py) Also, remove some commented-out code that is no longer needed. Co-authored-by: Mark Feldhousen --- tests/test_request_doc.py | 43 +++++++++++++++++++++++++++++++++++---- 1 file changed, 39 insertions(+), 4 deletions(-) diff --git a/tests/test_request_doc.py b/tests/test_request_doc.py index 4adbb3a..aa2f072 100644 --- a/tests/test_request_doc.py +++ b/tests/test_request_doc.py @@ -1,8 +1,14 @@ """Test RequestDoc model functionality.""" +# Standard Python Libraries +from datetime import time + +# Third-Party Libraries +import pytest + # cisagov Libraries from cyhy_db.models import RequestDoc -from cyhy_db.models.request_doc import Agency +from cyhy_db.models.request_doc import Agency, Window async def test_init(): @@ -22,6 +28,35 @@ async def test_init(): ), "id was not correctly set to agency acronym" -# @given(st.builds(RequestDoc)) -# def test_dump_model(instance): -# print(instance) +def test_parse_time_valid_time_str(): + """Test the parse_time validator with valid string input.""" + valid_time_str = "12:34:56" + parsed_time = Window.parse_time(valid_time_str) + assert parsed_time == time(12, 34, 56), "Failed to parse valid time string" + + +def test_parse_time_invalid_time_str(): + """Test the parse_time validator with invalid string input.""" + invalid_time_str = "invalid_time" + with pytest.raises( + ValueError, + match="time data 'invalid_time' does not match format '%H:%M:%S'", + ): + Window.parse_time(invalid_time_str) + + +def test_parse_time_valid_time_obj(): + """Test the parse_time validator with valid time input.""" + valid_time_obj = time(12, 34, 56) + parsed_time = Window.parse_time(valid_time_obj) + assert parsed_time == valid_time_obj, "Failed to parse valid time object" + + +def test_parse_time_invalid_type(): + """Test the parse_time validator with an invalid input type.""" + invalid_time_type = 12345 + with pytest.raises( + ValueError, + match="Invalid time format. Expected a string in '%H:%M:%S' format or datetime.time instance.", + ): + Window.parse_time(invalid_time_type) From 22ed790dbd38abe714061e191a66b3bb74898a82 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 15:27:21 -0400 Subject: [PATCH 094/139] Avoid using the reserved "_id" field name in TallyDoc Co-authored-by: Mark Feldhousen --- src/cyhy_db/models/tally_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/tally_doc.py b/src/cyhy_db/models/tally_doc.py index d1a69af..201df8a 100644 --- a/src/cyhy_db/models/tally_doc.py +++ b/src/cyhy_db/models/tally_doc.py @@ -38,7 +38,7 @@ class TallyDoc(Document): model_config = ConfigDict(extra="forbid") - _id: str # owner_id + id: str # owner_id counts: Counts = Field(default_factory=Counts) last_change: datetime = Field(default_factory=utcnow) From 6fbef453d42f70f6f407144c715e5036cca0e7b8 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 15:27:51 -0400 Subject: [PATCH 095/139] Add unit tests for TallyDoc Co-authored-by: Mark Feldhousen --- tests/test_tally_doc.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 tests/test_tally_doc.py diff --git a/tests/test_tally_doc.py b/tests/test_tally_doc.py new file mode 100644 index 0000000..b573772 --- /dev/null +++ b/tests/test_tally_doc.py @@ -0,0 +1,29 @@ +"""Test TallyDoc model functionality.""" + +# Standard Python Libraries +from datetime import datetime + +# cisagov Libraries +from cyhy_db.models.tally_doc import Counts, TallyDoc + + +async def test_tally_doc_creation(): + """Test TallyDoc creation.""" + tally_doc = TallyDoc(id="TALLY-TEST-1") + await tally_doc.insert() + fetched_doc = await TallyDoc.get(tally_doc.id) + assert fetched_doc is not None + assert fetched_doc.id == "TALLY-TEST-1" + assert fetched_doc.counts == Counts() + assert isinstance(fetched_doc.last_change, datetime) + + +async def test_tally_doc_last_change(): + """Test TallyDoc last_change update.""" + tally_doc = TallyDoc(id="TALLY-TEST-2") + await tally_doc.save() + initial_last_change = tally_doc.last_change + + # Save TallyDoc again to force the last_change timestamp to update + await tally_doc.save() + assert tally_doc.last_change > initial_last_change From 058a0c5f1b9183772284ee236fa54abdf3f657f5 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 17:07:19 -0400 Subject: [PATCH 096/139] Output mongodb URI at end of pytest run when used with --mongo-express option Co-authored-by: Mark Feldhousen --- tests/conftest.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 418f235..a1197c3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -91,7 +91,7 @@ def mongodb_container(docker_client, mongo_image_tag): @pytest.fixture(autouse=True, scope="session") -def mongo_express_container(docker_client, request): +def mongo_express_container(docker_client, db_uri, request): """Fixture for the Mongo Express test container.""" if not request.config.getoption("--mongo-express"): yield None @@ -114,7 +114,10 @@ def mongo_express_container(docker_client, request): def fin(): if request.config.getoption("--mongo-express"): print( - f"\n\nMongo Express is running at http://admin:pass@localhost:{MONGO_EXPRESS_PORT}" + f'\n\nMongoDB is accessible at {db_uri} with database named "{DATABASE_NAME}"' + ) + print( + f"Mongo Express is accessible at http://admin:pass@localhost:{MONGO_EXPRESS_PORT}\n" ) input("Press Enter to stop Mongo Express and MongoDB containers...") mongo_express_container.stop() From d3e4e03b93630968768bb4ecd5c113387486f224 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 17:23:19 -0400 Subject: [PATCH 097/139] Update README from skeleton to include useful information Co-authored-by: Mark Feldhousen --- README.md | 115 +++++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 101 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index 139655b..9bb85bd 100644 --- a/README.md +++ b/README.md @@ -5,20 +5,107 @@ [![Coverage Status](https://coveralls.io/repos/github/cisagov/cyhy-db/badge.svg?branch=develop)](https://coveralls.io/github/cisagov/cyhy-db?branch=develop) [![Known Vulnerabilities](https://snyk.io/test/github/cisagov/cyhy-db/develop/badge.svg)](https://snyk.io/test/github/cisagov/cyhy-db) -This is a generic skeleton project that can be used to quickly get a -new [cisagov](https://github.com/cisagov) Python library GitHub -project started. This skeleton project contains [licensing -information](LICENSE), as well as -[pre-commit hooks](https://pre-commit.com) and -[GitHub Actions](https://github.com/features/actions) configurations -appropriate for a Python library project. - -## New Repositories from a Skeleton ## - -Please see our [Project Setup guide](https://github.com/cisagov/development-guide/tree/develop/project_setup) -for step-by-step instructions on how to start a new repository from -a skeleton. This will save you time and effort when configuring a -new repository! +This repository implements a Python module for interacting with a Cyber Hygiene database. + +## Pre-requisites ## + +- [Python 3.10](https://www.python.org/downloads/) or newer +- A running [MongoDB](https://www.mongodb.com/) instance that you have access to + +## Starting a Local MongoDB Instance for Testing ## + +> [!IMPORTANT] +> This requires [Docker](https://www.docker.com/) to be installed in +> order for this to work. + +You can start a local MongoDB instance in a container with the following +command: + +```console +pytest -vs --mongo-express +``` + +> [!NOTE] +> The command `pytest -vs --mongo-express` not only starts a local +> MongoDB instance, but also runs all the `cyhy-db` unit tests, which will +> create various collections and documents in the database. + +Sample output (trimmed to highlight the important parts): + +```console + +MongoDB is accessible at mongodb://mongoadmin:secret@localhost:32859 with database named "test" +Mongo Express is accessible at http://admin:pass@localhost:8081 + +Press Enter to stop Mongo Express and MongoDB containers... +``` + +Note that the previous command will execute all of the `cyhy-db` unit tests, +which will create a variety of collections and documents. + +Based on the example output above, you can access the MongoDB instance at +`mongodb://mongoadmin:secret@localhost:32859` and the Mongo Express web +interface at `http://admin:pass@localhost:8081`. Note that the MongoDB +containers will remain running until you press "Enter" in that terminal. + +## Example Usage ## + +Once you have a MongoDB instance running, the sample Python code below +demonstrates how to initialize the database, create a new request document, save +it, and then retrieve it. + +```python +import asyncio +from cyhy_db import initialize_db +from cyhy_db.models import RequestDoc +from cyhy_db.models.request_doc import Agency + +async def main(): + # Initialize the CyHy database + await initialize_db("mongodb://mongoadmin:secret@localhost:32859", "test") + + # Create a new CyHy request document and save it in the database + new_request = RequestDoc( + agency=Agency(name="Acme Industries", acronym="AI") + ) + await new_request.save() + + # Find the request document and print its agency information + request = await RequestDoc.get("AI") + print(request.agency) + +asyncio.run(main()) +``` + +Output: + +```console +name='Acme Industries' acronym='AI' type=None contacts=[] location=None +``` + +## Additional Testing Options ## + +> [!WARNING] +> The default usernames and passwords are for testing purposes only. +> Do not use them in production environments. Always set strong, unique +> credentials. + +### Environment Variables ### + +| Variable | Description | Default | +|----------|-------------|---------| +| `MONGO_INITDB_ROOT_USERNAME` | The MongoDB root username | `mongoadmin` | +| `MONGO_INITDB_ROOT_PASSWORD` | The MongoDB root password | `secret` | +| `DATABASE_NAME` | The name of the database to use for testing | `test` | +| `MONGO_EXPRESS_PORT` | The port to use for the Mongo Express web interface | `8081` | + +### Pytest Options ### + +| Option | Description | Default | +|--------|-------------|---------| +| `--mongo-express` | Start a local MongoDB instance and Mongo Express web interface | n/a | +| `--mongo-image-tag` | The tag of the MongoDB Docker image to use | `docker.io/mongo:latest` | +| `--runslow` | Run slow tests | n/a | ## Contributing ## From c2157df6d95bb80063d701e09534252ba03a5bf7 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 17:31:43 -0400 Subject: [PATCH 098/139] Remove a TODO that is no longer relevant Co-authored-by: Mark Feldhousen --- src/cyhy_db/models/scan_doc.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 89173fc..07edd9a 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -74,7 +74,6 @@ async def reset_latest_flag_by_ip( ): """Reset the latest flag for all scans for a given IP address.""" if isinstance(ips, Iterable): - # TODO Figure out why coverage thinks this next line can exit early ip_ints = [int(ip_address(x)) for x in ips] else: ip_ints = [int(ip_address(ips))] From 7fec3813f6b0ee4cb4b7760180f4e3faaf0c26d6 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 17:44:52 -0400 Subject: [PATCH 099/139] Make all children of ScanDoc (HostScanDoc, PortScanDoc, VulnScanDoc) include the same indexes as ScanDoc Co-authored-by: Mark Feldhousen --- src/cyhy_db/models/host_scan_doc.py | 2 +- src/cyhy_db/models/port_scan_doc.py | 2 +- src/cyhy_db/models/vuln_scan_doc.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cyhy_db/models/host_scan_doc.py b/src/cyhy_db/models/host_scan_doc.py index ea2a1ad..892bac0 100644 --- a/src/cyhy_db/models/host_scan_doc.py +++ b/src/cyhy_db/models/host_scan_doc.py @@ -24,7 +24,7 @@ class Settings: """Beanie settings.""" name = "host_scans" - indexes = [ + indexes = ScanDoc.Settings.indexes + [ IndexModel( [("latest", ASCENDING), ("owner", ASCENDING)], name="latest_owner" ), diff --git a/src/cyhy_db/models/port_scan_doc.py b/src/cyhy_db/models/port_scan_doc.py index 2236f1b..ce66c3d 100644 --- a/src/cyhy_db/models/port_scan_doc.py +++ b/src/cyhy_db/models/port_scan_doc.py @@ -25,7 +25,7 @@ class Settings: """Beanie settings.""" name = "port_scans" - indexes = [ + indexes = ScanDoc.Settings.indexes + [ IndexModel( [("latest", ASCENDING), ("owner", ASCENDING), ("state", ASCENDING)], name="latest_owner_state", diff --git a/src/cyhy_db/models/vuln_scan_doc.py b/src/cyhy_db/models/vuln_scan_doc.py index 9060a5e..3552952 100644 --- a/src/cyhy_db/models/vuln_scan_doc.py +++ b/src/cyhy_db/models/vuln_scan_doc.py @@ -38,7 +38,7 @@ class Settings: """Beanie settings.""" name = "vuln_scans" - indexes = [ + indexes = ScanDoc.Settings.indexes + [ IndexModel( [("owner", ASCENDING), ("latest", ASCENDING), ("severity", ASCENDING)], name="owner_latest_severity", From aca07b4d047320ebff385505361feb9edc65cbab Mon Sep 17 00:00:00 2001 From: David Redmin Date: Wed, 2 Oct 2024 17:59:56 -0400 Subject: [PATCH 100/139] Update a comment Co-authored-by: Mark Feldhousen --- src/cyhy_db/models/scan_doc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 07edd9a..6b7bc8f 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -17,7 +17,8 @@ from .snapshot_doc import SnapshotDoc -class ScanDoc(Document): # TODO: Make this a BaseModel +# TODO: Figure out how to make this an abstract base class. BaseModel? +class ScanDoc(Document): """The scan document model.""" # Validate on assignment so ip_int is recalculated as ip is set From 9e06b475c00f0513402a78622fa29e0cfd99ba50 Mon Sep 17 00:00:00 2001 From: Felddy Date: Thu, 3 Oct 2024 11:25:51 -0400 Subject: [PATCH 101/139] Remove duplicate note text --- README.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/README.md b/README.md index 9bb85bd..c7d27f8 100644 --- a/README.md +++ b/README.md @@ -40,9 +40,6 @@ Mongo Express is accessible at http://admin:pass@localhost:8081 Press Enter to stop Mongo Express and MongoDB containers... ``` -Note that the previous command will execute all of the `cyhy-db` unit tests, -which will create a variety of collections and documents. - Based on the example output above, you can access the MongoDB instance at `mongodb://mongoadmin:secret@localhost:32859` and the Mongo Express web interface at `http://admin:pass@localhost:8081`. Note that the MongoDB From 7f4b9ab22f90b85c56469662e58de5ddcbf502bc Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 4 Oct 2024 13:54:50 -0400 Subject: [PATCH 102/139] Make ScanDoc abstract and document --- src/cyhy_db/db.py | 7 +++++-- src/cyhy_db/models/host_scan_doc.py | 2 +- src/cyhy_db/models/port_scan_doc.py | 2 +- src/cyhy_db/models/scan_doc.py | 23 ++++++++++++++++++----- src/cyhy_db/models/vuln_scan_doc.py | 2 +- tests/test_scan_doc.py | 6 +++++- 6 files changed, 31 insertions(+), 11 deletions(-) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index 6213484..ef416d5 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -14,10 +14,10 @@ PortScanDoc, ReportDoc, RequestDoc, - ScanDoc, SnapshotDoc, SystemControlDoc, TallyDoc, + TicketDoc, VulnScanDoc, ) @@ -31,13 +31,16 @@ PortScanDoc, RequestDoc, ReportDoc, - ScanDoc, SnapshotDoc, SystemControlDoc, TallyDoc, + TicketDoc, VulnScanDoc, ] +# Note: ScanDoc is intentionally excluded from the list of models to be imported +# or initialized because it is an abstract base class. + async def initialize_db(db_uri: str, db_name: str) -> AsyncIOMotorDatabase: """Initialize the database.""" diff --git a/src/cyhy_db/models/host_scan_doc.py b/src/cyhy_db/models/host_scan_doc.py index 892bac0..0bb9131 100644 --- a/src/cyhy_db/models/host_scan_doc.py +++ b/src/cyhy_db/models/host_scan_doc.py @@ -24,7 +24,7 @@ class Settings: """Beanie settings.""" name = "host_scans" - indexes = ScanDoc.Settings.indexes + [ + indexes = ScanDoc.Abstract_Settings.indexes + [ IndexModel( [("latest", ASCENDING), ("owner", ASCENDING)], name="latest_owner" ), diff --git a/src/cyhy_db/models/port_scan_doc.py b/src/cyhy_db/models/port_scan_doc.py index ce66c3d..f66fbb8 100644 --- a/src/cyhy_db/models/port_scan_doc.py +++ b/src/cyhy_db/models/port_scan_doc.py @@ -25,7 +25,7 @@ class Settings: """Beanie settings.""" name = "port_scans" - indexes = ScanDoc.Settings.indexes + [ + indexes = ScanDoc.Abstract_Settings.indexes + [ IndexModel( [("latest", ASCENDING), ("owner", ASCENDING), ("state", ASCENDING)], name="latest_owner_state", diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index 6b7bc8f..ddf8510 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -1,6 +1,7 @@ """ScanDoc model for use as the base of other scan document classes.""" # Standard Python Libraries +from abc import ABC from datetime import datetime from ipaddress import IPv4Address, ip_address from typing import Any, Dict, Iterable, List, Union @@ -17,9 +18,8 @@ from .snapshot_doc import SnapshotDoc -# TODO: Figure out how to make this an abstract base class. BaseModel? -class ScanDoc(Document): - """The scan document model.""" +class ScanDoc(Document, ABC): + """The abstract base class for scan-like documents.""" # Validate on assignment so ip_int is recalculated as ip is set model_config = ConfigDict(extra="forbid", validate_assignment=True) @@ -40,9 +40,22 @@ def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: return values class Settings: - """Beanie settings.""" + """Beanie settings to be used during testing.""" + + # These settings are intended for use only during testing. See + # Abstract_Settings below. + + name = "PyTest_ScanDocs" + + class Abstract_Settings: + """Beanie settings to be inherited by subclasses.""" + + # This class is intentionally not named "Settings" to prevent Beanie from + # automatically applying these indices, which would result in the creation of + # an unnecessary collection. Instead, subclasses should define their own + # "Settings" class and include these indices along with any additional + # subclass-specific indices. - name = "scandocs" indexes = [ IndexModel( [("latest", ASCENDING), ("ip_int", ASCENDING)], name="latest_ip" diff --git a/src/cyhy_db/models/vuln_scan_doc.py b/src/cyhy_db/models/vuln_scan_doc.py index 3552952..a48ebdd 100644 --- a/src/cyhy_db/models/vuln_scan_doc.py +++ b/src/cyhy_db/models/vuln_scan_doc.py @@ -38,7 +38,7 @@ class Settings: """Beanie settings.""" name = "vuln_scans" - indexes = ScanDoc.Settings.indexes + [ + indexes = ScanDoc.Abstract_Settings.indexes + [ IndexModel( [("owner", ASCENDING), ("latest", ASCENDING), ("severity", ASCENDING)], name="owner_latest_severity", diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index c605134..79d80e6 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -1,4 +1,4 @@ -"""Test ScanDoc model functionality.""" +"""Test ScanDoc abstract base class model functionality.""" # Standard Python Libraries import ipaddress @@ -16,6 +16,10 @@ VALID_IP_1_INT = int(ipaddress.ip_address(VALID_IP_1_STR)) VALID_IP_2_INT = int(ipaddress.ip_address(VALID_IP_2_STR)) +# Note: Running these tests will create a "ScanDoc" collection in the database. +# This collection is typically not created in a production environment since +# ScanDoc is an abstract base class. + def test_ip_int_init(): """Test IP address integer conversion on initialization. From 6f31267902639eab0ad0a385dbf1e1b76d56154b Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 4 Oct 2024 14:09:19 -0400 Subject: [PATCH 103/139] Add stub of TicketDoc model --- src/cyhy_db/models/__init__.py | 2 ++ src/cyhy_db/models/ticket_doc.py | 52 ++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) create mode 100644 src/cyhy_db/models/ticket_doc.py diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py index 7e2c947..6a5d339 100644 --- a/src/cyhy_db/models/__init__.py +++ b/src/cyhy_db/models/__init__.py @@ -25,6 +25,7 @@ from .request_doc import RequestDoc from .system_control_doc import SystemControlDoc from .tally_doc import TallyDoc +from .ticket_doc import TicketDoc __all__ = [ @@ -41,5 +42,6 @@ "SnapshotDoc", "SystemControlDoc", "TallyDoc", + "TicketDoc", "VulnScanDoc", ] diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py new file mode 100644 index 0000000..ef9802e --- /dev/null +++ b/src/cyhy_db/models/ticket_doc.py @@ -0,0 +1,52 @@ +"""The model for CyHy ticket documents.""" + +# Third-Party Libraries +from beanie import Document +from pymongo import ASCENDING, IndexModel + + +class TicketDoc(Document): + """The ticket document model.""" + + pass + + class Settings: + """Beanie settings.""" + + name = "tickets" + + indexes = [ + IndexModel( + [ + ("ip_int", ASCENDING), + ("port", ASCENDING), + ("protocol", ASCENDING), + ("source", ASCENDING), + ("source_id", ASCENDING), + ("open", ASCENDING), + ("false_positive", ASCENDING), + ], + name="ip_port_protocol_source_open_false_positive", + ), + IndexModel( + [("ip_int", ASCENDING), ("open", ASCENDING)], + name="ip_open", + ), + IndexModel( + [("open", ASCENDING), ("owner", ASCENDING)], + name="open_owner", + ), + IndexModel( + [("time_opened", ASCENDING), ("open", ASCENDING)], + name="time_opened", + ), + IndexModel( + [("last_change", ASCENDING)], + name="last_change", + ), + IndexModel( + [("time_closed", ASCENDING)], + name="time_closed", + sparse=True, + ), + ] From 3c04616ead57e4c01740880a6d4a74b978930220 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 4 Oct 2024 18:05:11 -0400 Subject: [PATCH 104/139] Improve comment for protocol enum --- src/cyhy_db/models/enum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index 270d9bb..3457f04 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -56,7 +56,7 @@ class PocType(Enum): class Protocol(Enum): - """Protocols.""" + """Network protocols.""" TCP = "tcp" UDP = "udp" From 74f7fc808e3832e077964144500a3ea630f5b2c9 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 4 Oct 2024 18:05:39 -0400 Subject: [PATCH 105/139] Rename TicketEvent --- src/cyhy_db/models/enum.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index 3457f04..b62726b 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -114,8 +114,8 @@ class Status(Enum): WAITING = "WAITING" -class TicketEvent(Enum): - """Ticket events.""" +class TicketAction(Enum): + """Actions for ticket events.""" CHANGED = "CHANGED" CLOSED = "CLOSED" From f2284baf77ceb91542f7c2f9dc22a80abf96a65a Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 4 Oct 2024 18:19:46 -0400 Subject: [PATCH 106/139] Add custom CyHy exceptions Co-authored-by: David Redmin --- src/cyhy_db/models/exceptions.py | 43 ++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 src/cyhy_db/models/exceptions.py diff --git a/src/cyhy_db/models/exceptions.py b/src/cyhy_db/models/exceptions.py new file mode 100644 index 0000000..6f2d4fd --- /dev/null +++ b/src/cyhy_db/models/exceptions.py @@ -0,0 +1,43 @@ +"""The exceptions used in CyHy.""" + + +class PortScanNotFoundException(Exception): + """Exception raised when a referenced PortScanDoc is not found.""" + + def __init__(self, ticket_id, port_scan_id, port_scan_time, *args): + """Initialize the exception with the given ticket ID, port scan ID, and port scan time. + + Args: + ticket_id (str): The ID of the ticket. + port_scan_id (str): The ID of the port scan. + port_scan_time (datetime): The time of the port scan. + *args: Additional arguments to pass to the base Exception class. + """ + message = "Ticket {}: referenced PortScanDoc {} at time {} not found".format( + ticket_id, port_scan_id, port_scan_time + ) + self.ticket_id = ticket_id + self.port_scan_id = port_scan_id + self.port_scan_time = port_scan_time + super().__init__(message, *args) + + +class VulnScanNotFoundException(Exception): + """Exception raised when a referenced VulnScanDoc is not found.""" + + def __init__(self, ticket_id, vuln_scan_id, vuln_scan_time, *args): + """Initialize the exception with the given ticket ID, vulnerability scan ID, and vulnerability scan time. + + Args: + ticket_id (str): The ID of the ticket. + vuln_scan_id (str): The ID of the vulnerability scan document. + vuln_scan_time (str): The time of the vulnerability scan. + *args: Additional arguments to pass to the base exception class. + """ + message = "Ticket {}: referenced VulnScanDoc {} at time {} not found".format( + ticket_id, vuln_scan_id, vuln_scan_time + ) + self.ticket_id = ticket_id + self.vuln_scan_id = vuln_scan_id + self.vuln_scan_time = vuln_scan_time + super().__init__(message, *args) From f772efd5482ea60d3a98a7b99d3f52bef5589322 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 4 Oct 2024 18:20:07 -0400 Subject: [PATCH 107/139] Add TIcketDoc implementation Co-authored-by: David Redmin --- src/cyhy_db/models/ticket_doc.py | 213 ++++++++++++++++++++++++++++++- 1 file changed, 211 insertions(+), 2 deletions(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index ef9802e..ecb2698 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -1,14 +1,72 @@ """The model for CyHy ticket documents.""" +# Standard Python Libraries +from datetime import datetime, timedelta +from ipaddress import IPv4Address +from typing import List, Optional, Tuple + # Third-Party Libraries -from beanie import Document +from beanie import ( + BeanieObjectId, + Document, + Insert, + Link, + Replace, + ValidateOnSave, + before_event, +) +from beanie.operators import In, Pull, Push +from pydantic import BaseModel, ConfigDict, Field from pymongo import ASCENDING, IndexModel +# cisagov Libraries +from cyhy_db.utils.time import utcnow + +from . import PortScanDoc, SnapshotDoc, VulnScanDoc +from .enum import Protocol, TicketAction +from .exceptions import PortScanNotFoundException, VulnScanNotFoundException + + +class EventDelta(BaseModel): + """The event delta model.""" + + from_: Optional[bool | float | int | str] = Field(..., alias="from") + key: str = Field(...) + to: Optional[bool | float | int | str] = Field(...) + + +class TicketEvent(BaseModel): + """The ticket event model.""" + + action: TicketAction + delta: Optional[EventDelta] = Field(default_factory=EventDelta()) + reason: str = Field(...) + reference: BeanieObjectId = Field(...) + time: datetime + class TicketDoc(Document): """The ticket document model.""" - pass + model_config = ConfigDict(extra="forbid") + + details: dict = Field(default_factory=dict) + events: list[TicketEvent] = Field(default_factory=list) + false_positive: bool = Field(default=False) + fp_expiration_date: Optional[datetime] = Field(default=None) + ip_int: int = Field(...) + ip: IPv4Address = Field(...) + last_change: datetime = Field(default_factory=utcnow) + loc: Optional[Tuple[float, float]] = Field(default=None) + open: bool = Field(default=True) + owner: str = Field(...) + port: int = Field(...) + protocol: Protocol = Field(...) + snapshots: Optional[List[Link[SnapshotDoc]]] = Field(default_factory=list) + source_id: int = Field(...) + source: str = Field(...) + time_closed: Optional[datetime] = Field(default=None) + time_opened: datetime = Field(default_factory=utcnow) class Settings: """Beanie settings.""" @@ -50,3 +108,154 @@ class Settings: sparse=True, ), ] + + @before_event(Insert, Replace, ValidateOnSave) + async def before_save(self): + """Do a false positive sanity check and set data just prior to saving a ticket document.""" + if self.false_positive and not self.open: + raise Exception("A ticket marked as a false positive cannot be closed.") + self.last_change = utcnow() + + def add_event(self, action, reason, reference=None, time=None, delta=None): + """Add an event to the list of ticket events.""" + if action not in TicketAction: + raise Exception( + 'Invalid action "' + action + '" cannot be added to ticket events.' + ) + if not time: + time = utcnow() + event = TicketEvent( + action=action, reason=reason, reference=reference, time=time + ) + if delta: + event.delta = delta + self["events"].append(event) + + def false_positive_dates(self): + """Return most recent false positive effective and expiration dates (if any).""" + if self.false_positive: + for event in reversed(self.events): + if not event.delta: + continue + if ( + event.action == TicketAction.CHANGED + and event.delta.key == "false_positive" + ): + return (event.time, self.fp_expiration_date) + return None + + def last_detection_date(self): + """Return date of most recent detection of a ticket's finding.""" + for event in reversed(self.events): + if event.action in [ + TicketAction.OPENED, + TicketAction.VERIFIED, + TicketAction.REOPENED, + ]: + return event.time + # This should never happen, but if we don't find any OPENED/VERIFIED/REOPENED events above, gracefully return time_opened + return self.time_opened + + async def latest_port(self): + """Return the last referenced port scan in the event list. + + This should only be used for tickets generated by portscans. + """ + for event in self.events[::-1]: + reference_id = event.get("reference") + if reference_id: + break + else: + raise Exception("No references found in ticket events:", self._id) + port = await PortScanDoc.get(reference_id) + if not port: + # This can occur when a port_scan has been archived + # Raise an exception with the info we have for this port_scan from the ticket + raise PortScanNotFoundException( + ticket_id=self._id, + port_scan_id=reference_id, + port_scan_time=event.time, + ) + return port + + async def latest_vuln(self): + """Return the last referenced vulnerability in the event list. + + This should only be used for tickets generated by vulnscans. + """ + for event in self.events[::-1]: + reference_id = event.get("reference") + if reference_id: + break + else: + raise Exception("No references found in ticket events:", self._id) + vuln = await VulnScanDoc.get(reference_id) + if not vuln: + # This can occur when a vuln_scan has been archived + # Raise an exception with the info we have for this vuln_scan from the ticket + raise VulnScanNotFoundException( + ticket_id=self._id, + vuln_scan_id=reference_id, + vuln_scan_time=event.time, + ) + return vuln + + async def set_false_positive(self, new_state: bool, reason: str, expire_days: int): + """Mark a ticket as a false positive.""" + if self.false_positive == new_state: + return + + # Define the event delta + delta = EventDelta( + from_=self.false_positive, to=new_state, key="false_positive" + ) + + # Update ticket state + self.false_positive = new_state + now = utcnow() + expiration_date = None + + if new_state: + # Only include the expiration date when setting false_positive to + # True + expiration_date = now + timedelta(days=expire_days) + + # If ticket is not open, re-open it; false positive tix should + # always be open + if not self.open: + self.open = True + self.time_closed = None + self.add_event( + action=TicketAction.REOPENED, + reason="setting false positive", + time=now, + ) + + # Add the change event + self.add_event( + action=TicketAction.CHANGED, reason=reason, time=now, delta=delta + ) + + # Set ticket expiration date if applicable + self.fp_expiration_date = expiration_date + + @classmethod + async def tag_open(cls, owners, snapshot_oid): + """Add a snapshot object ID to the snapshots field of all open tickets belonging to the specified owners.""" + await cls.find(cls.open is True, In(cls.owner, owners)).update_many( + Push({cls.snapshots: snapshot_oid}) + ) + + @classmethod + async def tag_matching(cls, existing_snapshot_oids, new_snapshot_oid): + """Add a new snapshot object ID to the snapshots field of all tickets whose snapshots field contain any of specified existing snapshot object IDs.""" + await cls.find(In(cls.snapshots, existing_snapshot_oids)).update_many( + Push({cls.snapshots: new_snapshot_oid}) + ) + + @classmethod + async def remove_tag(cls, snapshot_oid): + """Remove the specified snapshot object ID from the snapshots field of all tickets whose snapshots field contain that snapshot object ID.""" + await cls.find(In(cls.snapshots, snapshot_oid)).update_many( + Pull({cls.snapshots: snapshot_oid}) + ) From 8c45e19f7c465ae18a7af9bd0506e25d47cfd6b9 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 13:44:28 -0400 Subject: [PATCH 108/139] Use correct syntax when appending a TicketDoc event --- src/cyhy_db/models/ticket_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index ecb2698..d4e896a 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -129,7 +129,7 @@ def add_event(self, action, reason, reference=None, time=None, delta=None): ) if delta: event.delta = delta - self["events"].append(event) + self.events.append(event) def false_positive_dates(self): """Return most recent false positive effective and expiration dates (if any).""" From cf2b2d23f08a2490c2a4c8cf1c9ccaf11924e18e Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 13:45:25 -0400 Subject: [PATCH 109/139] Remove an unnecessary async declaration --- src/cyhy_db/models/ticket_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index d4e896a..16c09ef 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -200,7 +200,7 @@ async def latest_vuln(self): ) return vuln - async def set_false_positive(self, new_state: bool, reason: str, expire_days: int): + def set_false_positive(self, new_state: bool, reason: str, expire_days: int): """Mark a ticket as a false positive.""" if self.false_positive == new_state: return From 3e3d4709743b1aba8db0605dfd4e20e2e18fb730 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 13:47:25 -0400 Subject: [PATCH 110/139] Correctly define TicketEvent delta to default to None Since delta is optional, it makes sense to default to None. Co-authored-by: Mark Feldhousen --- src/cyhy_db/models/ticket_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index 16c09ef..2e42653 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -39,7 +39,7 @@ class TicketEvent(BaseModel): """The ticket event model.""" action: TicketAction - delta: Optional[EventDelta] = Field(default_factory=EventDelta()) + delta: Optional[EventDelta] = Field(default=None) reason: str = Field(...) reference: BeanieObjectId = Field(...) time: datetime From 50c8c33094a17893eb584d1b33d98e5b7c103d2f Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 13:49:04 -0400 Subject: [PATCH 111/139] Prefer serialization_alias to alias For more info, see https://docs.pydantic.dev/latest/concepts/fields/#field-aliases Co-authored-by: Mark Feldhousen --- src/cyhy_db/models/ticket_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index 2e42653..b64ef42 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -30,7 +30,7 @@ class EventDelta(BaseModel): """The event delta model.""" - from_: Optional[bool | float | int | str] = Field(..., alias="from") + from_: Optional[bool | float | int | str] = Field(..., serialization_alias="from") key: str = Field(...) to: Optional[bool | float | int | str] = Field(...) From 5786d37d5a9c2d30be2f0133d6c5bce4e8477857 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 17:22:07 -0400 Subject: [PATCH 112/139] Use correct query syntax --- src/cyhy_db/models/ticket_doc.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index b64ef42..d5bab0f 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -242,9 +242,11 @@ def set_false_positive(self, new_state: bool, reason: str, expire_days: int): @classmethod async def tag_open(cls, owners, snapshot_oid): """Add a snapshot object ID to the snapshots field of all open tickets belonging to the specified owners.""" - await cls.find(cls.open is True, In(cls.owner, owners)).update_many( - Push({cls.snapshots: snapshot_oid}) - ) + # flake8 E712 is "comparison to True should be 'if cond is True:' or 'if + # cond:'" but this is unavoidable due to Beanie syntax. + await cls.find( + cls.open == True, In(cls.owner, owners) # noqa E712 + ).update_many(Push({cls.snapshots: snapshot_oid})) @classmethod async def tag_matching(cls, existing_snapshot_oids, new_snapshot_oid): From ab431a6e1483c3921024577a3ca95c0b648c1b2e Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 17:23:01 -0400 Subject: [PATCH 113/139] The "In" query operator requires an array --- src/cyhy_db/models/ticket_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index d5bab0f..c6861ad 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -258,6 +258,6 @@ async def tag_matching(cls, existing_snapshot_oids, new_snapshot_oid): @classmethod async def remove_tag(cls, snapshot_oid): """Remove the specified snapshot object ID from the snapshots field of all tickets whose snapshots field contain that snapshot object ID.""" - await cls.find(In(cls.snapshots, snapshot_oid)).update_many( + await cls.find(In(cls.snapshots, [snapshot_oid])).update_many( Pull({cls.snapshots: snapshot_oid}) ) From 720abfea280e6f9798343602e80a7df3c1b41be8 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 20:33:58 -0400 Subject: [PATCH 114/139] Correct issues with latest_port and latest_vuln in TicketDoc --- src/cyhy_db/models/ticket_doc.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index c6861ad..4fef20a 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -162,17 +162,17 @@ async def latest_port(self): This should only be used for tickets generated by portscans. """ for event in self.events[::-1]: - reference_id = event.get("reference") + reference_id = event.reference if reference_id: break else: - raise Exception("No references found in ticket events:", self._id) + raise Exception("No references found in ticket events: " + str(self.id)) port = await PortScanDoc.get(reference_id) if not port: # This can occur when a port_scan has been archived # Raise an exception with the info we have for this port_scan from the ticket raise PortScanNotFoundException( - ticket_id=self._id, + ticket_id=self.id, port_scan_id=reference_id, port_scan_time=event.time, ) @@ -184,17 +184,17 @@ async def latest_vuln(self): This should only be used for tickets generated by vulnscans. """ for event in self.events[::-1]: - reference_id = event.get("reference") + reference_id = event.reference if reference_id: break else: - raise Exception("No references found in ticket events:", self._id) + raise Exception("No references found in ticket events: " + str(self.id)) vuln = await VulnScanDoc.get(reference_id) if not vuln: # This can occur when a vuln_scan has been archived # Raise an exception with the info we have for this vuln_scan from the ticket raise VulnScanNotFoundException( - ticket_id=self._id, + ticket_id=self.id, vuln_scan_id=reference_id, vuln_scan_time=event.time, ) From 6e579af6201379d96513ea19d209e1224140b55c Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 20:34:59 -0400 Subject: [PATCH 115/139] Make TicketEvent.reference optional, defaulting to None Some event types (e.g. CHANGED) do not require a reference. --- src/cyhy_db/models/ticket_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index 4fef20a..cbd2e8f 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -41,7 +41,7 @@ class TicketEvent(BaseModel): action: TicketAction delta: Optional[EventDelta] = Field(default=None) reason: str = Field(...) - reference: BeanieObjectId = Field(...) + reference: Optional[BeanieObjectId] = Field(default=None) time: datetime From f6668db93e44d5ad76781d15b7690019beec4a15 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 20:38:58 -0400 Subject: [PATCH 116/139] Add unit tests for TicketDoc --- tests/test_ticket_doc.py | 425 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 425 insertions(+) create mode 100644 tests/test_ticket_doc.py diff --git a/tests/test_ticket_doc.py b/tests/test_ticket_doc.py new file mode 100644 index 0000000..fc2678a --- /dev/null +++ b/tests/test_ticket_doc.py @@ -0,0 +1,425 @@ +"""Test TicketDoc model functionality.""" + +# Standard Python Libraries +from datetime import timedelta +from ipaddress import IPv4Address +from unittest.mock import AsyncMock, patch + +# Third-Party Libraries +from beanie import PydanticObjectId +import pytest + +# cisagov Libraries +from cyhy_db.models.enum import Protocol, TicketAction +from cyhy_db.models.exceptions import ( + PortScanNotFoundException, + VulnScanNotFoundException, +) +from cyhy_db.models.port_scan_doc import PortScanDoc +from cyhy_db.models.snapshot_doc import SnapshotDoc +from cyhy_db.models.ticket_doc import EventDelta, TicketDoc, TicketEvent +from cyhy_db.models.vuln_scan_doc import VulnScanDoc +from cyhy_db.utils.time import utcnow + +VALID_IP_1_STR = "0.0.0.1" +VALID_IP_1_INT = int(IPv4Address(VALID_IP_1_STR)) + + +def sample_ticket(): + """Create a sample TicketDoc object.""" + return TicketDoc( + ip_int=VALID_IP_1_INT, + ip=IPv4Address(VALID_IP_1_STR), + owner="TICKET-TEST-1", + port=80, + protocol=Protocol.TCP, + source_id=1, + source="test", + ) + + +def test_init(): + """Test TicketDoc object initialization.""" + # Create a TicketDoc object + ticket_doc = sample_ticket() + + # Verify that default values are set correctly + assert ticket_doc.details == {}, "details was not set to an empty dict" + assert ticket_doc.events == [], "events was not set to an empty list" + assert ticket_doc.false_positive is False, "false_positive was not set to False" + assert ticket_doc.last_change is not None, "last_change was not set" + assert ticket_doc.open is True, "open was not set to True" + assert ticket_doc.snapshots == [], "snapshots was not set to an empty list" + assert ticket_doc.time_closed is None, "time_closed was not set to None" + assert ticket_doc.time_opened is not None, "time_opened was not set" + + +async def test_save(): + """Test saving a TicketDoc object to the database.""" + # Create a TicketDoc object and save it to the DB + ticket_doc = sample_ticket() + await ticket_doc.save() + + # Find ticket in DB and confirm it was saved correctly + ticket_doc_db = await TicketDoc.find_one(TicketDoc.ip_int == VALID_IP_1_INT) + assert ticket_doc_db is not None, "ticket_doc was not saved to the database" + + +async def test_save_with_event(): + """Test saving a ticket that contains an event.""" + ticket_doc = await TicketDoc.find_one(TicketDoc.ip_int == VALID_IP_1_INT) + ticket_doc.set_false_positive( + new_state=True, reason="Test set false positive", expire_days=30 + ) + await ticket_doc.save() + + # Find ticket in DB and confirm it was saved correctly + ticket_doc_db = await TicketDoc.find_one(TicketDoc.ip_int == VALID_IP_1_INT) + assert ticket_doc_db is not None, "ticket_doc was not saved to the database" + assert ticket_doc_db.events[0].action == TicketAction.CHANGED + assert ticket_doc_db.events[0].delta.from_ is False + assert ticket_doc_db.events[0].delta.to is True + + +async def test_before_save(): + """Test the before_save method.""" + ticket_doc = sample_ticket() + ticket_doc.false_positive = True + ticket_doc.open = False + with pytest.raises( + Exception, match="A ticket marked as a false positive cannot be closed." + ): + await ticket_doc.save() + + +def test_add_event(): + """Test adding an event to a ticket.""" + ticket_doc = sample_ticket() + ticket_doc.add_event(action=TicketAction.OPENED, reason="Test reason") + assert len(ticket_doc.events) == 1, "event was not added to the ticket" + assert ticket_doc.events[0].action == TicketAction.OPENED + assert ticket_doc.events[0].reason == "Test reason" + + +def test_add_event_exception(): + """Test adding an invalid event to a ticket.""" + ticket_doc = sample_ticket() + with pytest.raises( + Exception, match='Invalid action "INVALID" cannot be added to ticket events.' + ): + ticket_doc.add_event(action="INVALID", reason="Test reason") + + +def test_set_false_positive_true(): + """Test setting a ticket as false positive.""" + ticket_doc = sample_ticket() + ticket_doc.set_false_positive( + new_state=True, reason="Test set false positive", expire_days=30 + ) + assert ticket_doc.false_positive is True, "ticket was not set as false positive" + assert ( + ticket_doc.fp_expiration_date is not None + ), "false positive expiration date was not set" + + +def test_set_false_positive_no_change(): + """Test setting a ticket that was already false positive to false positive.""" + ticket_doc = sample_ticket() + + ticket_doc.set_false_positive( + new_state=True, reason="Test set false positive", expire_days=30 + ) + fp_expiration_date = ticket_doc.fp_expiration_date + + ticket_doc.set_false_positive( + new_state=True, reason="Test set false positive again", expire_days=60 + ) + assert ( + ticket_doc.false_positive is True + ), "ticket should have remained a false positive" + assert ( + ticket_doc.fp_expiration_date == fp_expiration_date + ), "false positive expiration date should not have changed" + + +def test_set_false_positive_false(): + """Test setting a ticket as false positive false.""" + ticket_doc = sample_ticket() + ticket_doc.set_false_positive( + new_state=True, reason="Test set false positive true", expire_days=30 + ) + + assert ( + ticket_doc.fp_expiration_date is not None + ), "false positive expiration date was not set" + + ticket_doc.set_false_positive( + new_state=False, reason="Test set false positive false", expire_days=0 + ) + assert ( + ticket_doc.false_positive is False + ), "ticket should not still be false positive" + assert ( + ticket_doc.fp_expiration_date is None + ), "false positive expiration date was not cleared" + + +def test_set_false_positive_on_closed_ticket(): + """Test setting a closed ticket as false positive.""" + ticket_doc = sample_ticket() + # Close the ticket + ticket_doc.open = False + ticket_doc.time_closed = utcnow() + + ticket_doc.set_false_positive( + expire_days=30, + new_state=True, + reason="Test set false positive on closed ticket", + ) + assert ticket_doc.open is True, "ticket should have been reopened" + assert ticket_doc.false_positive is True, "ticket should be false positive" + assert ticket_doc.time_closed is None, "ticket should not have a time_closed" + assert ticket_doc.events[-2].action == TicketAction.REOPENED + assert ticket_doc.events[-2].reason == "setting false positive" + assert ticket_doc.events[-2].time is not None + + +def test_false_positive_dates(): + """Test getting false positive dates.""" + ticket_doc = sample_ticket() + + # Set ticket as false positive + ticket_doc.set_false_positive( + new_state=True, reason="Test set false positive", expire_days=30 + ) + fp_dates = ticket_doc.false_positive_dates() + assert fp_dates is not None + + # Add another sample event + ticket_doc.add_event( + action=TicketAction.UNVERIFIED, reason="Test reason", time=utcnow() + ) + assert ticket_doc.false_positive_dates() == fp_dates + + # Unset ticket as false positive + ticket_doc.false_positive = False + event = TicketEvent( + action=TicketAction.CHANGED, + delta=EventDelta(from_=True, to=False, key="false_positive"), + reason="Test false positive expired", + reference=None, + time=utcnow(), + ) + ticket_doc.events.append(event) + assert ticket_doc.false_positive_dates() is None + + +def test_false_positive_dates_edge_cases(): + """Test getting false positive dates edge cases.""" + ticket_doc = sample_ticket() + ticket_doc.false_positive = True + assert ticket_doc.false_positive_dates() is None + + # Add a sample non-false-positive CHANGED event + test_delta = EventDelta(from_=False, to=True, key="test_key") + ticket_doc.add_event( + action=TicketAction.CHANGED, + delta=test_delta, + reason="Test reason", + time=utcnow(), + ) + assert ticket_doc.false_positive_dates() is None + + +def test_last_detection_date(): + """Test getting the last detection date.""" + ticket_doc = sample_ticket() + ticket_doc.add_event( + action=TicketAction.OPENED, reason="Test reason", time=utcnow() + ) + detection_date = ticket_doc.last_detection_date() + assert detection_date == ticket_doc.events[0].time + + +def test_last_detection_date_edge_case(): + """Test an edge case of last_detection_date.""" + ticket_doc = sample_ticket() + ticket_doc.add_event( + action=TicketAction.CLOSED, reason="Test reason", time=utcnow() + ) + detection_date = ticket_doc.last_detection_date() + assert detection_date == ticket_doc.time_opened + + +async def test_tagging(): + """Test tag_open, tag_matching, and remove_tag.""" + # Find our test ticket in the DB + ticket_doc_db = await TicketDoc.find_one(TicketDoc.ip_int == VALID_IP_1_INT) + test_owner = ticket_doc_db.owner + assert len(ticket_doc_db.snapshots) == 0 + + # Create a test snapshot and save it to the DB + snapshot_end_time = utcnow() + snapshot_start_time = snapshot_end_time - timedelta(days=1) + test_snapshot_1 = SnapshotDoc( + owner=test_owner, end_time=snapshot_end_time, start_time=snapshot_start_time + ) + await test_snapshot_1.save() + assert test_snapshot_1 not in ticket_doc_db.snapshots + + # Use tag_open() to tag the ticket with the snapshot ID + await TicketDoc.tag_open(owners=[test_owner], snapshot_oid=test_snapshot_1.id) + + updated_ticket = await TicketDoc.find_one(TicketDoc.ip_int == VALID_IP_1_INT) + # I'm not using fetch_links=True in the find_one() above because I can't get + # it to work correctly. Instead, I'm using fetch_all_links() below. + await updated_ticket.fetch_all_links() + assert len(updated_ticket.snapshots) == 1 + assert test_snapshot_1 in updated_ticket.snapshots + + # Create another test snapshot and save it to the DB + snapshot_end_time = utcnow() + snapshot_start_time = snapshot_end_time - timedelta(days=1) + test_snapshot_2 = SnapshotDoc( + owner=test_owner, end_time=snapshot_end_time, start_time=snapshot_start_time + ) + await test_snapshot_2.save() + assert test_snapshot_2 not in updated_ticket.snapshots + + # Use tag_matching() to tag the ticket with the new snapshot ID + await TicketDoc.tag_matching( + existing_snapshot_oids=[test_snapshot_1.id], + new_snapshot_oid=test_snapshot_2.id, + ) + + updated_ticket = await TicketDoc.find_one(TicketDoc.ip_int == VALID_IP_1_INT) + # I'm not using fetch_links=True in the find_one() above because I can't get + # it to work correctly. Instead, I'm using fetch_all_links() below. + await updated_ticket.fetch_all_links() + assert len(updated_ticket.snapshots) == 2 + assert test_snapshot_2 in updated_ticket.snapshots + + # Use remove_tag() to remove the test_snapshot_2.id from the ticket + await TicketDoc.remove_tag(snapshot_oid=test_snapshot_2.id) + + updated_ticket = await TicketDoc.find_one(TicketDoc.ip_int == VALID_IP_1_INT) + # I'm not using fetch_links=True in the find_one() above because I can't get + # it to work correctly. Instead, I'm using fetch_all_links() below. + await updated_ticket.fetch_all_links() + assert len(updated_ticket.snapshots) == 1 + assert test_snapshot_2 not in updated_ticket.snapshots + + +async def test_latest_port(): + """Test the latest_port method.""" + ticket_doc = sample_ticket() + ticket_doc.id = PydanticObjectId() + reference_id = PydanticObjectId() + # Add an event with our test reference ID + ticket_doc.add_event( + action=TicketAction.OPENED, + reason="Test reason", + reference=reference_id, + time=utcnow(), + ) + # Add another event without a reference ID + ticket_doc.add_event( + action=TicketAction.VERIFIED, + reason="Test reason", + time=utcnow(), + ) + + # Create a dummy port scan document with the reference ID + mock_doc = AsyncMock() + mock_doc.id = reference_id + + with patch.object(PortScanDoc, "get", return_value=mock_doc): + port = await ticket_doc.latest_port() + assert ( + port.id == reference_id + ), "latest_port did not return the correct port scan" + + +async def test_latest_port_no_references(): + """Test the latest_port method when there are no references.""" + ticket_doc = sample_ticket() + ticket_doc.id = PydanticObjectId() + + with pytest.raises( + Exception, match=("No references found in ticket events: " + str(ticket_doc.id)) + ): + await ticket_doc.latest_port() + + +async def test_latest_port_not_found(): + """Test the latest_port method when the port scan is not found.""" + ticket_doc = sample_ticket() + reference_id = PydanticObjectId() + ticket_doc.add_event( + action=TicketAction.OPENED, + reason="Test reason", + reference=reference_id, + time=utcnow(), + ) + + with pytest.raises(PortScanNotFoundException): + # Mock PortScanDoc.get to return None + with patch.object(PortScanDoc, "get", return_value=None): + await ticket_doc.latest_port() + + +async def test_latest_vuln(): + """Test the latest_vuln method.""" + ticket_doc = sample_ticket() + reference_id = PydanticObjectId() + # Add an event with our test reference ID + ticket_doc.add_event( + action=TicketAction.OPENED, + reason="Test reason", + reference=reference_id, + time=utcnow(), + ) + # Add another event without a reference ID + ticket_doc.add_event( + action=TicketAction.VERIFIED, + reason="Test reason", + time=utcnow(), + ) + + # Create a dummy port scan document with the reference ID + mock_doc = AsyncMock() + mock_doc._id = reference_id + + with patch.object(VulnScanDoc, "get", return_value=mock_doc): + vuln = await ticket_doc.latest_vuln() + assert ( + vuln._id == reference_id + ), "latest_vuln did not return the correct vuln scan" + + +async def test_latest_vuln_no_references(): + """Test the latest_vuln method when there are no references.""" + ticket_doc = sample_ticket() + ticket_doc.id = PydanticObjectId() + + with pytest.raises( + Exception, match=("No references found in ticket events: " + str(ticket_doc.id)) + ): + await ticket_doc.latest_vuln() + + +async def test_latest_vuln_not_found(): + """Test the latest_vuln method when the port scan is not found.""" + ticket_doc = sample_ticket() + reference_id = PydanticObjectId() + ticket_doc.add_event( + action=TicketAction.OPENED, + reason="Test reason", + reference=reference_id, + time=utcnow(), + ) + + with pytest.raises(VulnScanNotFoundException): + # Mock VulnScanDoc.get to return None + with patch.object(VulnScanDoc, "get", return_value=None): + await ticket_doc.latest_vuln() From 41a2275cea9c07a6f2bb98b1aa91b251c9963a95 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Mon, 7 Oct 2024 20:58:01 -0400 Subject: [PATCH 117/139] Make TicketDoc.add_event() raise the same exception in all currently-supported Python versions --- src/cyhy_db/models/ticket_doc.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index cbd2e8f..fcc1c98 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -118,7 +118,11 @@ async def before_save(self): def add_event(self, action, reason, reference=None, time=None, delta=None): """Add an event to the list of ticket events.""" - if action not in TicketAction: + try: + action = TicketAction(action) + # If action is not in the enumerated TicketAction class, Python 3.10 + # and 3.11 throw a TypeError, while Python 3.12 throws a ValueError + except (TypeError, ValueError): raise Exception( 'Invalid action "' + action + '" cannot be added to ticket events.' ) From 32806099f2f3e9c0fcde66cafacca1e82da0d590 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Tue, 8 Oct 2024 15:45:06 -0400 Subject: [PATCH 118/139] Remove commented-out legacy functions There are a lot of functions in the Python 2 version of HostDoc (such as the ones being deleted here) that may or may not be used. Instead of porting them all over to cyhy-db, we plan to port them as they are needed and rewrite them to take advantage of Python 3. --- src/cyhy_db/models/host_doc.py | 33 --------------------------------- 1 file changed, 33 deletions(-) diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 6989d55..3dc7106 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -124,41 +124,8 @@ def set_state(self, nmap_says_up, has_open_ports, reason=None): # Instead of porting them all over, we should just port them as they are needed. # And rewrite things that can be done better in Python 3. - # @classmethod - # async def get_count(cls, owner: str, stage: Stage, status: Status): - # return await cls.count( - # cls.owner == owner, cls.stage == stage, cls.status == status - # ) - @classmethod @deprecated("Use HostDoc.find_one(HostDoc.ip == ip) instead.") async def get_by_ip(cls, ip: IPv4Address): """Return a host document with the given IP address.""" return await cls.find_one(cls.ip == ip) - - # @classmethod - # @deprecated("Use HostDoc.find_one(HostDoc.ip == ip).owner instead.") - # async def get_owner_of_ip(cls, ip: IPv4Address): - # host = await cls.get_by_ip(ip) - # return host.owner - - # @classmethod - # async def get_some_for_stage( - # cls, - # stage: Stage, - # count: int, - # owner: Optional[str] = None, - # waiting: bool = False, - # ): - # if waiting: - # status = {"$in": [Status.READY, Status.WAITING]} - # else: - # status = Status.READY - - # query = cls.find(cls.status == status, cls.stage == stage) - # if owner is not None: - # query = query.find(cls.owner == owner) - - # # Sorting and limiting the results - # results = await query.sort([("priority", 1), ("r", 1)]).limit(count).to_list() - # return results From 88987f64a7de2e157026c702ea9597d7ad219602 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 8 Oct 2024 15:58:34 -0400 Subject: [PATCH 119/139] Convert enums to StrEnum and use auto() This requires a minimum Python version of 3.11 --- .github/workflows/build.yml | 3 - README.md | 2 +- setup.py | 3 +- src/cyhy_db/models/enum.py | 132 +++++++++++++++---------------- src/cyhy_db/models/ticket_doc.py | 4 +- 5 files changed, 70 insertions(+), 74 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 5cc0632..b60d2b6 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -180,7 +180,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.10" - "3.11" - "3.12" steps: @@ -283,7 +282,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.10" - "3.11" - "3.12" steps: @@ -335,7 +333,6 @@ jobs: fail-fast: false matrix: python-version: - - "3.10" - "3.11" - "3.12" steps: diff --git a/README.md b/README.md index c7d27f8..c3afd47 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ This repository implements a Python module for interacting with a Cyber Hygiene ## Pre-requisites ## -- [Python 3.10](https://www.python.org/downloads/) or newer +- [Python 3.11](https://www.python.org/downloads/) or newer - A running [MongoDB](https://www.mongodb.com/) instance that you have access to ## Starting a Local MongoDB Instance for Testing ## diff --git a/setup.py b/setup.py index 7c01781..0071d36 100644 --- a/setup.py +++ b/setup.py @@ -75,12 +75,11 @@ def get_version(version_file): # that you indicate whether you support Python 2, Python 3 or both. "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: Implementation :: CPython", ], - python_requires=">=3.10", + python_requires=">=3.11", # What does your project relate to? keywords=["cyhy", "database"], packages=find_packages(where="src"), diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index b62726b..09c324f 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -1,125 +1,125 @@ """The enumerations used in CyHy.""" # Standard Python Libraries -from enum import Enum +from enum import StrEnum, auto -class AgencyType(Enum): +class AgencyType(StrEnum): """Agency types.""" - FEDERAL = "FEDERAL" - LOCAL = "LOCAL" - PRIVATE = "PRIVATE" - STATE = "STATE" - TERRITORIAL = "TERRITORIAL" - TRIBAL = "TRIBAL" + FEDERAL = auto() + LOCAL = auto() + PRIVATE = auto() + STATE = auto() + TERRITORIAL = auto() + TRIBAL = auto() -class ControlAction(Enum): +class ControlAction(StrEnum): """Commander control actions.""" - PAUSE = "PAUSE" - STOP = "STOP" + PAUSE = auto() + STOP = auto() -class ControlTarget(Enum): +class ControlTarget(StrEnum): """Commander control targets.""" - COMMANDER = "COMMANDER" + COMMANDER = auto() -class CVSSVersion(Enum): +class CVSSVersion(StrEnum): """CVSS versions.""" - V2 = "2.0" - V3 = "3.0" - V3_1 = "3.1" + V2 = auto() + V3 = auto() + V3_1 = auto() -class DayOfWeek(Enum): +class DayOfWeek(StrEnum): """Days of the week.""" - MONDAY = "MONDAY" - TUESDAY = "TUESDAY" - WEDNESDAY = "WEDNESDAY" - THURSDAY = "THURSDAY" - FRIDAY = "FRIDAY" - SATURDAY = "SATURDAY" - SUNDAY = "SUNDAY" + MONDAY = auto() + TUESDAY = auto() + WEDNESDAY = auto() + THURSDAY = auto() + FRIDAY = auto() + SATURDAY = auto() + SUNDAY = auto() -class PocType(Enum): +class PocType(StrEnum): """Point of contact types.""" - DISTRO = "DISTRO" - TECHNICAL = "TECHNICAL" + DISTRO = auto() + TECHNICAL = auto() -class Protocol(Enum): +class Protocol(StrEnum): """Network protocols.""" - TCP = "tcp" - UDP = "udp" + TCP = auto() + UDP = auto() -class ReportPeriod(Enum): +class ReportPeriod(StrEnum): """CyHy reporting periods.""" - MONTHLY = "MONTHLY" - QUARTERLY = "QUARTERLY" - WEEKLY = "WEEKLY" + MONTHLY = auto() + QUARTERLY = auto() + WEEKLY = auto() -class ReportType(Enum): +class ReportType(StrEnum): """CyHy report types.""" - BOD = "BOD" - CYBEX = "CYBEX" - CYHY = "CYHY" - CYHY_THIRD_PARTY = "CYHY_THIRD_PARTY" - DNSSEC = "DNSSEC" - PHISHING = "PHISHING" + BOD = auto() + CYBEX = auto() + CYHY = auto() + CYHY_THIRD_PARTY = auto() + DNSSEC = auto() + PHISHING = auto() -class ScanType(Enum): +class ScanType(StrEnum): """CyHy scan types.""" - CYHY = "CYHY" - DNSSEC = "DNSSEC" - PHISHING = "PHISHING" + CYHY = auto() + DNSSEC = auto() + PHISHING = auto() -class Scheduler(Enum): +class Scheduler(StrEnum): """CyHy schedulers.""" - PERSISTENT1 = "PERSISTENT1" + PERSISTENT1 = auto() -class Stage(Enum): +class Stage(StrEnum): """CyHy scan stages.""" - BASESCAN = "BASESCAN" # TODO: Delete if unused - NETSCAN1 = "NETSCAN1" - NETSCAN2 = "NETSCAN2" - PORTSCAN = "PORTSCAN" - VULNSCAN = "VULNSCAN" + BASESCAN = auto() # TODO: Delete if unused + NETSCAN1 = auto() + NETSCAN2 = auto() + PORTSCAN = auto() + VULNSCAN = auto() -class Status(Enum): +class Status(StrEnum): """CyHy scan statuses.""" - DONE = "DONE" - READY = "READY" - RUNNING = "RUNNING" - WAITING = "WAITING" + DONE = auto() + READY = auto() + RUNNING = auto() + WAITING = auto() -class TicketAction(Enum): +class TicketAction(StrEnum): """Actions for ticket events.""" - CHANGED = "CHANGED" - CLOSED = "CLOSED" - OPENED = "OPENED" - REOPENED = "REOPENED" - UNVERIFIED = "UNVERIFIED" - VERIFIED = "VERIFIED" + CHANGED = auto() + CLOSED = auto() + OPENED = auto() + REOPENED = auto() + UNVERIFIED = auto() + VERIFIED = auto() diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index fcc1c98..a4d7deb 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -120,8 +120,8 @@ def add_event(self, action, reason, reference=None, time=None, delta=None): """Add an event to the list of ticket events.""" try: action = TicketAction(action) - # If action is not in the enumerated TicketAction class, Python 3.10 - # and 3.11 throw a TypeError, while Python 3.12 throws a ValueError + # If action is not in the enumerated TicketAction class, Python 3.11 + # throws a TypeError, while Python 3.12 throws a ValueError except (TypeError, ValueError): raise Exception( 'Invalid action "' + action + '" cannot be added to ticket events.' From fa69bfcbc94888e48ebe12591075199c830d5954 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 8 Oct 2024 16:01:03 -0400 Subject: [PATCH 120/139] Remove unused BASESCAN enum value in Stage Co-authored-by: David Redmin --- src/cyhy_db/models/enum.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index 09c324f..f41cbc6 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -98,7 +98,6 @@ class Scheduler(StrEnum): class Stage(StrEnum): """CyHy scan stages.""" - BASESCAN = auto() # TODO: Delete if unused NETSCAN1 = auto() NETSCAN2 = auto() PORTSCAN = auto() From f2210f3ef69cac0791c49c6f57fa3e8390f3fa15 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 8 Oct 2024 16:24:28 -0400 Subject: [PATCH 121/139] Reverting change to CVSSVersion enum --- src/cyhy_db/models/enum.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/cyhy_db/models/enum.py b/src/cyhy_db/models/enum.py index f41cbc6..02b77b0 100644 --- a/src/cyhy_db/models/enum.py +++ b/src/cyhy_db/models/enum.py @@ -31,9 +31,9 @@ class ControlTarget(StrEnum): class CVSSVersion(StrEnum): """CVSS versions.""" - V2 = auto() - V3 = auto() - V3_1 = auto() + V2 = "2.0" + V3 = "3.0" + V3_1 = "3.1" class DayOfWeek(StrEnum): From c1d1a29b9dfe626f15ab6285112d8855684dce65 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 8 Oct 2024 16:25:35 -0400 Subject: [PATCH 122/139] Convert test to use CVSS enumeration --- tests/test_cve.py | 29 +++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/tests/test_cve.py b/tests/test_cve.py index d9f999d..1ac525e 100644 --- a/tests/test_cve.py +++ b/tests/test_cve.py @@ -6,20 +6,21 @@ # cisagov Libraries from cyhy_db.models import CVE +from cyhy_db.models.enum import CVSSVersion severity_params = [ - ("2.0", 10, 4), - ("2.0", 7.0, 3), - ("2.0", 4.0, 2), - ("2.0", 0.0, 1), - ("3.0", 9.0, 4), - ("3.0", 7.0, 3), - ("3.0", 4.0, 2), - ("3.0", 0.0, 1), - ("3.1", 9.0, 4), - ("3.1", 7.0, 3), - ("3.1", 4.0, 2), - ("3.1", 0.0, 1), + (CVSSVersion.V2, 10, 4), + (CVSSVersion.V2, 7.0, 3), + (CVSSVersion.V2, 4.0, 2), + (CVSSVersion.V2, 0.0, 1), + (CVSSVersion.V3, 9.0, 4), + (CVSSVersion.V3, 7.0, 3), + (CVSSVersion.V3, 4.0, 2), + (CVSSVersion.V3, 0.0, 1), + (CVSSVersion.V3_1, 9.0, 4), + (CVSSVersion.V3_1, 7.0, 3), + (CVSSVersion.V3_1, 4.0, 2), + (CVSSVersion.V3_1, 0.0, 1), ] @@ -36,12 +37,12 @@ def test_calculate_severity(version, score, expected_severity): def test_invalid_cvss_score(bad_score): """Test that an invalid CVSS score raises a ValueError.""" with pytest.raises(ValidationError): - CVE(cvss_version="3.1", cvss_score=bad_score, id="test-cve") + CVE(cvss_version=CVSSVersion.V3_1, cvss_score=bad_score, id="test-cve") async def test_save(): """Test that the severity is calculated correctly on save.""" - cve = CVE(cvss_version="3.1", cvss_score=9.0, id="test-cve") + cve = CVE(cvss_version=CVSSVersion.V3_1, cvss_score=9.0, id="test-cve") await cve.save() # Saving the object saved_cve = await CVE.get("test-cve") # Retrieving the object From 4613a1e22d1f5fb863c0962d81b85e086707c744 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 8 Oct 2024 16:33:31 -0400 Subject: [PATCH 123/139] Rename CVE to CVEDoc This change normalizes the names of all document classes. Co-authored-by: David Redmin --- src/cyhy_db/db.py | 4 ++-- src/cyhy_db/models/__init__.py | 4 ++-- src/cyhy_db/models/{cve.py => cve_doc.py} | 2 +- tests/test_connection.py | 4 ++-- tests/{test_cve.py => test_cve_doc.py} | 10 +++++----- tests/test_data_generator.py | 4 ++-- 6 files changed, 14 insertions(+), 14 deletions(-) rename src/cyhy_db/models/{cve.py => cve_doc.py} (98%) rename tests/{test_cve.py => test_cve_doc.py} (79%) diff --git a/src/cyhy_db/db.py b/src/cyhy_db/db.py index ef416d5..c74e14b 100644 --- a/src/cyhy_db/db.py +++ b/src/cyhy_db/db.py @@ -5,7 +5,7 @@ from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase from .models import ( - CVE, + CVEDoc, HostDoc, HostScanDoc, KEVDoc, @@ -22,7 +22,7 @@ ) ALL_MODELS: list[type[Document] | type[View] | str] = [ - CVE, + CVEDoc, HostDoc, HostScanDoc, KEVDoc, diff --git a/src/cyhy_db/models/__init__.py b/src/cyhy_db/models/__init__.py index 6a5d339..8e06369 100644 --- a/src/cyhy_db/models/__init__.py +++ b/src/cyhy_db/models/__init__.py @@ -17,7 +17,7 @@ from .report_doc import ReportDoc # Other documents -from .cve import CVE +from .cve_doc import CVEDoc from .host_doc import HostDoc from .kev_doc import KEVDoc from .notification_doc import NotificationDoc @@ -29,7 +29,7 @@ __all__ = [ - "CVE", + "CVEDoc", "HostDoc", "HostScanDoc", "KEVDoc", diff --git a/src/cyhy_db/models/cve.py b/src/cyhy_db/models/cve_doc.py similarity index 98% rename from src/cyhy_db/models/cve.py rename to src/cyhy_db/models/cve_doc.py index 1c65370..9003c81 100644 --- a/src/cyhy_db/models/cve.py +++ b/src/cyhy_db/models/cve_doc.py @@ -10,7 +10,7 @@ from .enum import CVSSVersion -class CVE(Document): +class CVEDoc(Document): """The CVE document model.""" # Validate on assignment so severity is calculated diff --git a/tests/test_connection.py b/tests/test_connection.py index 4febd58..3ddda07 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -9,7 +9,7 @@ # cisagov Libraries from cyhy_db.db import ALL_MODELS, initialize_db -from cyhy_db.models import CVE +from cyhy_db.models import CVEDoc async def test_connection_motor(db_uri, db_name): @@ -23,7 +23,7 @@ async def test_connection_motor(db_uri, db_name): async def test_connection_beanie(): """Test a simple database query.""" # Attempt to find a document in the empty CVE collection - result = await CVE.get("CVE-2024-DOES-NOT-EXIST") + result = await CVEDoc.get("CVE-2024-DOES-NOT-EXIST") assert result is None, "Expected no document to be found" diff --git a/tests/test_cve.py b/tests/test_cve_doc.py similarity index 79% rename from tests/test_cve.py rename to tests/test_cve_doc.py index 1ac525e..5dfe5a9 100644 --- a/tests/test_cve.py +++ b/tests/test_cve_doc.py @@ -5,7 +5,7 @@ import pytest # cisagov Libraries -from cyhy_db.models import CVE +from cyhy_db.models import CVEDoc from cyhy_db.models.enum import CVSSVersion severity_params = [ @@ -27,7 +27,7 @@ @pytest.mark.parametrize("version, score, expected_severity", severity_params) def test_calculate_severity(version, score, expected_severity): """Test that the severity is calculated correctly.""" - cve = CVE(id="CVE-2024-0128", cvss_version=version, cvss_score=score) + cve = CVEDoc(id="CVE-2024-0128", cvss_version=version, cvss_score=score) assert ( cve.severity == expected_severity ), f"Failed for CVSS {version} with score {score}" @@ -37,14 +37,14 @@ def test_calculate_severity(version, score, expected_severity): def test_invalid_cvss_score(bad_score): """Test that an invalid CVSS score raises a ValueError.""" with pytest.raises(ValidationError): - CVE(cvss_version=CVSSVersion.V3_1, cvss_score=bad_score, id="test-cve") + CVEDoc(cvss_version=CVSSVersion.V3_1, cvss_score=bad_score, id="test-cve") async def test_save(): """Test that the severity is calculated correctly on save.""" - cve = CVE(cvss_version=CVSSVersion.V3_1, cvss_score=9.0, id="test-cve") + cve = CVEDoc(cvss_version=CVSSVersion.V3_1, cvss_score=9.0, id="test-cve") await cve.save() # Saving the object - saved_cve = await CVE.get("test-cve") # Retrieving the object + saved_cve = await CVEDoc.get("test-cve") # Retrieving the object assert saved_cve is not None, "CVE not saved correctly" assert saved_cve.severity == 4, "Severity not calculated correctly on save" diff --git a/tests/test_data_generator.py b/tests/test_data_generator.py index 9edffe9..a61604a 100644 --- a/tests/test_data_generator.py +++ b/tests/test_data_generator.py @@ -19,7 +19,7 @@ from pytest_factoryboy import register # cisagov Libraries -from cyhy_db.models import CVE, RequestDoc +from cyhy_db.models import CVEDoc, RequestDoc from cyhy_db.models.enum import ( AgencyType, CVSSVersion, @@ -86,7 +86,7 @@ class CVEFactory(factory.Factory): class Meta: """Meta class for CVEFactory.""" - model = CVE + model = CVEDoc id = factory.LazyFunction(lambda: generic.cyhy_provider.cve_id()) # The following lines generate warnings from bandit about "Standard From 1c534615935438f96775c9a0b729c770a7939c41 Mon Sep 17 00:00:00 2001 From: David Redmin Date: Fri, 11 Oct 2024 14:52:13 -0400 Subject: [PATCH 124/139] Sort attributes in all document classes --- src/cyhy_db/models/cve_doc.py | 4 +-- src/cyhy_db/models/host_doc.py | 12 ++++---- src/cyhy_db/models/host_scan_doc.py | 4 +-- src/cyhy_db/models/notification_doc.py | 4 +-- src/cyhy_db/models/place_doc.py | 24 +++++++-------- src/cyhy_db/models/port_scan_doc.py | 5 ++-- src/cyhy_db/models/report_doc.py | 4 +-- src/cyhy_db/models/request_doc.py | 2 +- src/cyhy_db/models/snapshot_doc.py | 38 ++++++++++++------------ src/cyhy_db/models/system_control_doc.py | 4 +-- src/cyhy_db/models/tally_doc.py | 2 +- src/cyhy_db/models/vuln_scan_doc.py | 6 ++-- 12 files changed, 55 insertions(+), 54 deletions(-) diff --git a/src/cyhy_db/models/cve_doc.py b/src/cyhy_db/models/cve_doc.py index 9003c81..f581d7f 100644 --- a/src/cyhy_db/models/cve_doc.py +++ b/src/cyhy_db/models/cve_doc.py @@ -16,10 +16,10 @@ class CVEDoc(Document): # Validate on assignment so severity is calculated model_config = ConfigDict(extra="forbid", validate_assignment=True) - # CVE ID as a string - id: str = Indexed(primary_field=True) # type: ignore[assignment] cvss_score: float = Field(ge=0.0, le=10.0) cvss_version: CVSSVersion = Field(default=CVSSVersion.V3_1) + # CVE ID as a string + id: str = Indexed(primary_field=True) # type: ignore[assignment] severity: int = Field(ge=1, le=4, default=1) @model_validator(mode="before") diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 3dc7106..4a25f53 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -30,16 +30,16 @@ class HostDoc(Document): # IP address as an integer id: int = Field(default_factory=int) # type: ignore[assignment] ip: IPv4Address = Field(...) - owner: str = Field(...) last_change: datetime = Field(default_factory=utcnow) - next_scan: Optional[datetime] = Field(default=None) - state: State = Field(default_factory=lambda: State(reason="new", up=False)) - stage: Stage = Field(default=Stage.NETSCAN1) - status: Status = Field(default=Status.WAITING) + latest_scan: Dict[Stage, datetime] = Field(default_factory=dict) loc: Optional[Tuple[float, float]] = Field(default=None) + next_scan: Optional[datetime] = Field(default=None) + owner: str = Field(...) priority: int = Field(default=0) r: float = Field(default_factory=random.random) - latest_scan: Dict[Stage, datetime] = Field(default_factory=dict) + stage: Stage = Field(default=Stage.NETSCAN1) + state: State = Field(default_factory=lambda: State(reason="new", up=False)) + status: Status = Field(default=Status.WAITING) @model_validator(mode="before") def calculate_ip_int(cls, values: Dict[str, Any]) -> Dict[str, Any]: diff --git a/src/cyhy_db/models/host_scan_doc.py b/src/cyhy_db/models/host_scan_doc.py index 0bb9131..4d5dc93 100644 --- a/src/cyhy_db/models/host_scan_doc.py +++ b/src/cyhy_db/models/host_scan_doc.py @@ -15,10 +15,10 @@ class HostScanDoc(ScanDoc): model_config = ConfigDict(extra="forbid") - name: str accuracy: int - line: int classes: List[dict] = [] + line: int + name: str class Settings: """Beanie settings.""" diff --git a/src/cyhy_db/models/notification_doc.py b/src/cyhy_db/models/notification_doc.py index 1261395..9754e71 100644 --- a/src/cyhy_db/models/notification_doc.py +++ b/src/cyhy_db/models/notification_doc.py @@ -13,11 +13,11 @@ class NotificationDoc(Document): model_config = ConfigDict(extra="forbid") - ticket_id: BeanieObjectId = Field(...) # ticket id that triggered the notification - ticket_owner: str # owner of the ticket generated_for: List[str] = Field( default=[] ) # list of owners built as notifications are generated + ticket_id: BeanieObjectId = Field(...) # ticket id that triggered the notification + ticket_owner: str # owner of the ticket class Settings: """Beanie settings.""" diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index e220a85..6ad48dc 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -13,23 +13,23 @@ class PlaceDoc(Document): model_config = ConfigDict(extra="forbid") + clazz: str = Field(alias="class") # 'class' is a reserved keyword in Python + country_name: str + country: str + county_fips: Optional[str] = None + county: Optional[str] = None + elevation_feet: Optional[int] = None + elevation_meters: Optional[int] = None # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html id: int = Field(default_factory=int) # type: ignore[assignment] + latitude_dec: float + latitude_dms: Optional[str] = None + longitude_dec: float + longitude_dms: Optional[str] = None name: str - clazz: str = Field(alias="class") # 'class' is a reserved keyword in Python - state: str state_fips: str state_name: str - county: Optional[str] = None - county_fips: Optional[str] = None - country: str - country_name: str - latitude_dms: Optional[str] = None - longitude_dms: Optional[str] = None - latitude_dec: float - longitude_dec: float - elevation_meters: Optional[int] = None - elevation_feet: Optional[int] = None + state: str class Settings: """Beanie settings.""" diff --git a/src/cyhy_db/models/port_scan_doc.py b/src/cyhy_db/models/port_scan_doc.py index f66fbb8..bbf1304 100644 --- a/src/cyhy_db/models/port_scan_doc.py +++ b/src/cyhy_db/models/port_scan_doc.py @@ -15,11 +15,12 @@ class PortScanDoc(ScanDoc): """The port scan document model.""" model_config = ConfigDict(extra="forbid") - protocol: Protocol + port: int + protocol: Protocol + reason: str service: Dict = {} # Assuming no specific structure for "service" state: str - reason: str class Settings: """Beanie settings.""" diff --git a/src/cyhy_db/models/report_doc.py b/src/cyhy_db/models/report_doc.py index 4526707..2ce4f88 100644 --- a/src/cyhy_db/models/report_doc.py +++ b/src/cyhy_db/models/report_doc.py @@ -19,10 +19,10 @@ class ReportDoc(Document): model_config = ConfigDict(extra="forbid") - owner: str generated_time: datetime = Field(default_factory=utcnow) - snapshots: List[Link[SnapshotDoc]] + owner: str report_types: List[ReportType] + snapshots: List[Link[SnapshotDoc]] class Settings: """Beanie settings.""" diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index e729723..4874ace 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -101,10 +101,10 @@ class RequestDoc(Document): model_config = ConfigDict(extra="forbid") - id: str = Field(default=BOGUS_ID) # type: ignore[assignment] agency: Agency children: List[Link["RequestDoc"]] = Field(default=[]) enrolled: datetime = Field(default_factory=utcnow) + id: str = Field(default=BOGUS_ID) # type: ignore[assignment] init_stage: Stage = Field(default=Stage.NETSCAN1) key: Optional[str] = Field(default=None) networks: List[IPv4Network] = Field(default=[]) diff --git a/src/cyhy_db/models/snapshot_doc.py b/src/cyhy_db/models/snapshot_doc.py index dddc431..d4bffd2 100644 --- a/src/cyhy_db/models/snapshot_doc.py +++ b/src/cyhy_db/models/snapshot_doc.py @@ -30,14 +30,14 @@ class WorldData(BaseModel): model_config = ConfigDict(extra="forbid") + cvss_average_all: float = 0.0 + cvss_average_vulnerable: float = 0.0 host_count: int = 0 - vulnerable_host_count: int = 0 - vulnerabilities: VulnerabilityCounts = Field(default_factory=VulnerabilityCounts) unique_vulnerabilities: VulnerabilityCounts = Field( default_factory=VulnerabilityCounts ) - cvss_average_all: float = 0.0 - cvss_average_vulnerable: float = 0.0 + vulnerable_host_count: int = 0 + vulnerabilities: VulnerabilityCounts = Field(default_factory=VulnerabilityCounts) class TicketMetrics(BaseModel): @@ -45,8 +45,8 @@ class TicketMetrics(BaseModel): model_config = ConfigDict(extra="forbid") - median: int = 0 max: int = 0 + median: int = 0 class TicketOpenMetrics(BaseModel): @@ -80,29 +80,29 @@ class SnapshotDoc(Document): model_config = ConfigDict(extra="forbid") - owner: str = Field(...) + addresses_scanned: int = Field(default=0) + cvss_average_all: float = Field(default=0.0) + cvss_average_vulnerable: float = Field(default=0.0) descendants_included: List[str] = Field(default=[]) - last_change: datetime = Field(default_factory=utcnow) - start_time: datetime = Field(...) end_time: datetime = Field(...) + host_count: int = Field(default=0) + last_change: datetime = Field(default_factory=utcnow) latest: bool = Field(default=True) + networks: List[IPv4Network] = Field(default=[]) + owner: str = Field(...) port_count: int = Field(default=0) - unique_port_count: int = Field(default=0) + services: Dict = Field(default_factory=dict) + start_time: datetime = Field(...) + tix_msec_open: TicketOpenMetrics = Field(default_factory=TicketOpenMetrics) + tix_msec_to_close: TicketCloseMetrics = Field(default_factory=TicketCloseMetrics) unique_operating_systems: int = Field(default=0) - host_count: int = Field(default=0) - vulnerable_host_count: int = Field(default=0) - vulnerabilities: VulnerabilityCounts = Field(default_factory=VulnerabilityCounts) + unique_port_count: int = Field(default=0) unique_vulnerabilities: VulnerabilityCounts = Field( default_factory=VulnerabilityCounts ) - cvss_average_all: float = Field(default=0.0) - cvss_average_vulnerable: float = Field(default=0.0) + vulnerabilities: VulnerabilityCounts = Field(default_factory=VulnerabilityCounts) + vulnerable_host_count: int = Field(default=0) world: WorldData = Field(default_factory=WorldData) - networks: List[IPv4Network] = Field(default=[]) - addresses_scanned: int = Field(default=0) - services: Dict = Field(default_factory=dict) - tix_msec_open: TicketOpenMetrics = Field(default_factory=TicketOpenMetrics) - tix_msec_to_close: TicketCloseMetrics = Field(default_factory=TicketCloseMetrics) class Settings: """Beanie settings.""" diff --git a/src/cyhy_db/models/system_control_doc.py b/src/cyhy_db/models/system_control_doc.py index fd0714e..6486bdb 100644 --- a/src/cyhy_db/models/system_control_doc.py +++ b/src/cyhy_db/models/system_control_doc.py @@ -21,11 +21,11 @@ class SystemControlDoc(Document): model_config = ConfigDict(extra="forbid") action: ControlAction + completed: bool = False # Set to True when after the action has occurred + reason: str # Free-form, for UI / Logging sender: str # Free-form, for UI / Logging target: ControlTarget - reason: str # Free-form, for UI / Logging time: datetime = Field(default_factory=utcnow) # creation time - completed: bool = False # Set to True when after the action has occurred class Settings: """Beanie settings.""" diff --git a/src/cyhy_db/models/tally_doc.py b/src/cyhy_db/models/tally_doc.py index 201df8a..4c1a47f 100644 --- a/src/cyhy_db/models/tally_doc.py +++ b/src/cyhy_db/models/tally_doc.py @@ -38,8 +38,8 @@ class TallyDoc(Document): model_config = ConfigDict(extra="forbid") - id: str # owner_id counts: Counts = Field(default_factory=Counts) + id: str # owner_id last_change: datetime = Field(default_factory=utcnow) @before_event(Insert, Replace, ValidateOnSave) diff --git a/src/cyhy_db/models/vuln_scan_doc.py b/src/cyhy_db/models/vuln_scan_doc.py index a48ebdd..9df40c1 100644 --- a/src/cyhy_db/models/vuln_scan_doc.py +++ b/src/cyhy_db/models/vuln_scan_doc.py @@ -16,9 +16,6 @@ class VulnScanDoc(ScanDoc): model_config = ConfigDict(extra="forbid") - protocol: Protocol - port: int - service: str cvss_base_score: float cvss_vector: str description: str @@ -29,7 +26,10 @@ class VulnScanDoc(ScanDoc): plugin_name: str plugin_publication_date: datetime plugin_type: str + port: int + protocol: Protocol risk_factor: str + service: str severity: int solution: str synopsis: str From 8fff2b4a827029cb75dee35f4f4ebe3b1b47ad30 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 11 Oct 2024 17:09:12 -0400 Subject: [PATCH 125/139] Clean up dependencies - Remove unnecessary pins. - Remove unused dependencies. - `coveralls` comment is no longer required. - `setuptools` pin and comment is no longer required. See: - https://github.com/cisagov/cyhy-config/commit/8df2e692d2920743c40368a496ca154769b2985a --- setup.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/setup.py b/setup.py index 0071d36..a03d864 100644 --- a/setup.py +++ b/setup.py @@ -88,24 +88,16 @@ def get_version(version_file): py_modules=[splitext(basename(path))[0] for path in glob("src/*.py")], include_package_data=True, install_requires=[ - "beanie == 1.26.0", - "docopt == 0.6.2", + "beanie", "pydantic[email, hypothesis]", # hypothesis plugin is currently disabled: https://github.com/pydantic/pydantic/issues/4682 - "schema == 0.7.7", - "setuptools >= 73.0.1", + "setuptools", ], extras_require={ "test": [ "pytest-asyncio", "coverage", - # coveralls 1.11.0 added a service number for calls from - # GitHub Actions. This caused a regression which resulted in a 422 - # response from the coveralls API with the message: - # Unprocessable Entity for url: https://coveralls.io/api/v1/jobs - # 1.11.1 fixed this issue, but to ensure expected behavior we'll pin - # to never grab the regression version. - "coveralls != 1.11.0", - "docker == 7.1.0", + "coveralls", + "docker", "hypothesis", "mimesis-factory", "mimesis", From e5fdec21e6b9dbb9d44544c63dd5cc64d1bf4d7e Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 11 Oct 2024 17:11:03 -0400 Subject: [PATCH 126/139] Replace magic string with proper enum Co-authored-by: David Redmin --- src/cyhy_db/models/cve_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/cve_doc.py b/src/cyhy_db/models/cve_doc.py index f581d7f..36986fa 100644 --- a/src/cyhy_db/models/cve_doc.py +++ b/src/cyhy_db/models/cve_doc.py @@ -25,7 +25,7 @@ class CVEDoc(Document): @model_validator(mode="before") def calculate_severity(cls, values: Dict[str, Any]) -> Dict[str, Any]: """Calculate CVE severity based on the CVSS score and version.""" - if values["cvss_version"] == "2.0": + if values["cvss_version"] == CVSSVersion.V2: if values["cvss_score"] == 10: values["severity"] = 4 elif values["cvss_score"] >= 7.0: From 0f917aeb84c953c25cf412a0eabce627aca6f7b3 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 11 Oct 2024 17:21:46 -0400 Subject: [PATCH 127/139] Revert hard coding of Python version --- .github/workflows/build.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index b60d2b6..fb1d7d1 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -62,7 +62,7 @@ jobs: - id: setup-python uses: actions/setup-python@v5 with: - python-version: "3.12" + python-version: ${{ steps.setup-env.outputs.python-version }} # We need the Go version and Go cache location for the actions/cache step, # so the Go installation must happen before that. - id: setup-go @@ -244,7 +244,7 @@ jobs: - id: setup-python uses: actions/setup-python@v5 with: - python-version: "3.12" + python-version: ${{ steps.setup-env.outputs.python-version }} - uses: actions/cache@v3 env: BASE_CACHE_KEY: "${{ github.job }}-${{ runner.os }}-\ From 5b2bb40c9984243cc01a870175ff07ad05167925 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 15 Oct 2024 13:41:58 -0400 Subject: [PATCH 128/139] Work around type assignment error Add a mypy ignore directive to TallyDoc. Document all uses of this directive in a new issue. --- src/cyhy_db/models/host_doc.py | 1 + src/cyhy_db/models/kev_doc.py | 1 + src/cyhy_db/models/place_doc.py | 1 + src/cyhy_db/models/request_doc.py | 1 + src/cyhy_db/models/tally_doc.py | 4 +++- 5 files changed, 7 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 4a25f53..066e8e8 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -27,6 +27,7 @@ class HostDoc(Document): model_config = ConfigDict(extra="forbid") + # See: https://github.com/cisagov/cyhy-db/issues/7 # IP address as an integer id: int = Field(default_factory=int) # type: ignore[assignment] ip: IPv4Address = Field(...) diff --git a/src/cyhy_db/models/kev_doc.py b/src/cyhy_db/models/kev_doc.py index 4f9986d..f2745be 100644 --- a/src/cyhy_db/models/kev_doc.py +++ b/src/cyhy_db/models/kev_doc.py @@ -10,6 +10,7 @@ class KEVDoc(Document): model_config = ConfigDict(extra="forbid") + # See: https://github.com/cisagov/cyhy-db/issues/7 id: str = Field(default_factory=str) # type: ignore[assignment] known_ransomware: bool diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index 6ad48dc..125ba67 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -20,6 +20,7 @@ class PlaceDoc(Document): county: Optional[str] = None elevation_feet: Optional[int] = None elevation_meters: Optional[int] = None + # See: https://github.com/cisagov/cyhy-db/issues/7 # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html id: int = Field(default_factory=int) # type: ignore[assignment] latitude_dec: float diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 4874ace..60bff70 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -104,6 +104,7 @@ class RequestDoc(Document): agency: Agency children: List[Link["RequestDoc"]] = Field(default=[]) enrolled: datetime = Field(default_factory=utcnow) + # See: https://github.com/cisagov/cyhy-db/issues/7 id: str = Field(default=BOGUS_ID) # type: ignore[assignment] init_stage: Stage = Field(default=Stage.NETSCAN1) key: Optional[str] = Field(default=None) diff --git a/src/cyhy_db/models/tally_doc.py b/src/cyhy_db/models/tally_doc.py index 4c1a47f..54522b2 100644 --- a/src/cyhy_db/models/tally_doc.py +++ b/src/cyhy_db/models/tally_doc.py @@ -39,7 +39,9 @@ class TallyDoc(Document): model_config = ConfigDict(extra="forbid") counts: Counts = Field(default_factory=Counts) - id: str # owner_id + # See: https://github.com/cisagov/cyhy-db/issues/7 + # Owner ID string + id: str # type: ignore[assignment] last_change: datetime = Field(default_factory=utcnow) @before_event(Insert, Replace, ValidateOnSave) From fd82fe948ca88fe2e8f949842987468b3472d564 Mon Sep 17 00:00:00 2001 From: Felddy Date: Tue, 15 Oct 2024 15:55:38 -0400 Subject: [PATCH 129/139] Work around type assignment error Add a mypy ignore directive to TallyDoc. Document all uses of this directive in a new issue. --- src/cyhy_db/models/cve_doc.py | 1 + src/cyhy_db/models/host_doc.py | 1 + src/cyhy_db/models/kev_doc.py | 1 + src/cyhy_db/models/place_doc.py | 1 + src/cyhy_db/models/request_doc.py | 1 + src/cyhy_db/models/tally_doc.py | 4 +++- 6 files changed, 8 insertions(+), 1 deletion(-) diff --git a/src/cyhy_db/models/cve_doc.py b/src/cyhy_db/models/cve_doc.py index 36986fa..86956ae 100644 --- a/src/cyhy_db/models/cve_doc.py +++ b/src/cyhy_db/models/cve_doc.py @@ -18,6 +18,7 @@ class CVEDoc(Document): cvss_score: float = Field(ge=0.0, le=10.0) cvss_version: CVSSVersion = Field(default=CVSSVersion.V3_1) + # See: https://github.com/cisagov/cyhy-db/issues/7 # CVE ID as a string id: str = Indexed(primary_field=True) # type: ignore[assignment] severity: int = Field(ge=1, le=4, default=1) diff --git a/src/cyhy_db/models/host_doc.py b/src/cyhy_db/models/host_doc.py index 4a25f53..066e8e8 100644 --- a/src/cyhy_db/models/host_doc.py +++ b/src/cyhy_db/models/host_doc.py @@ -27,6 +27,7 @@ class HostDoc(Document): model_config = ConfigDict(extra="forbid") + # See: https://github.com/cisagov/cyhy-db/issues/7 # IP address as an integer id: int = Field(default_factory=int) # type: ignore[assignment] ip: IPv4Address = Field(...) diff --git a/src/cyhy_db/models/kev_doc.py b/src/cyhy_db/models/kev_doc.py index 4f9986d..f2745be 100644 --- a/src/cyhy_db/models/kev_doc.py +++ b/src/cyhy_db/models/kev_doc.py @@ -10,6 +10,7 @@ class KEVDoc(Document): model_config = ConfigDict(extra="forbid") + # See: https://github.com/cisagov/cyhy-db/issues/7 id: str = Field(default_factory=str) # type: ignore[assignment] known_ransomware: bool diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index 6ad48dc..125ba67 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -20,6 +20,7 @@ class PlaceDoc(Document): county: Optional[str] = None elevation_feet: Optional[int] = None elevation_meters: Optional[int] = None + # See: https://github.com/cisagov/cyhy-db/issues/7 # GNIS FEATURE_ID (INCITS 446-2008) - https://geonames.usgs.gov/domestic/index.html id: int = Field(default_factory=int) # type: ignore[assignment] latitude_dec: float diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 4874ace..60bff70 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -104,6 +104,7 @@ class RequestDoc(Document): agency: Agency children: List[Link["RequestDoc"]] = Field(default=[]) enrolled: datetime = Field(default_factory=utcnow) + # See: https://github.com/cisagov/cyhy-db/issues/7 id: str = Field(default=BOGUS_ID) # type: ignore[assignment] init_stage: Stage = Field(default=Stage.NETSCAN1) key: Optional[str] = Field(default=None) diff --git a/src/cyhy_db/models/tally_doc.py b/src/cyhy_db/models/tally_doc.py index 4c1a47f..54522b2 100644 --- a/src/cyhy_db/models/tally_doc.py +++ b/src/cyhy_db/models/tally_doc.py @@ -39,7 +39,9 @@ class TallyDoc(Document): model_config = ConfigDict(extra="forbid") counts: Counts = Field(default_factory=Counts) - id: str # owner_id + # See: https://github.com/cisagov/cyhy-db/issues/7 + # Owner ID string + id: str # type: ignore[assignment] last_change: datetime = Field(default_factory=utcnow) @before_event(Insert, Replace, ValidateOnSave) From 07724b23d80129fea518e373f5825469882bc800 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 18 Oct 2024 12:04:01 -0400 Subject: [PATCH 130/139] Rename reserved word deconfliction to match previously used pattern The reserved word `from` was replaced with `from_` in `TicketDoc`. This follows that pattern. Co-authored-by: David Redmin Co-authored-by: Jeremy Frasier --- src/cyhy_db/models/place_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index 125ba67..a170b0a 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -13,7 +13,7 @@ class PlaceDoc(Document): model_config = ConfigDict(extra="forbid") - clazz: str = Field(alias="class") # 'class' is a reserved keyword in Python + class_: str = Field(alias="class") # 'class' is a reserved keyword in Python country_name: str country: str county_fips: Optional[str] = None From 438c199f0ecd5616c3e7a444d772bc81f6a1ee0c Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 18 Oct 2024 12:14:27 -0400 Subject: [PATCH 131/139] Remove field alias This breaking change is documented in: - Data migration from `cyhy-core` database #8 Co-authored-by: David Redmin --- src/cyhy_db/models/request_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/request_doc.py b/src/cyhy_db/models/request_doc.py index 60bff70..39778a1 100644 --- a/src/cyhy_db/models/request_doc.py +++ b/src/cyhy_db/models/request_doc.py @@ -68,7 +68,7 @@ class ScanLimit(BaseModel): model_config = ConfigDict(extra="forbid") - scan_type: ScanType = Field(..., alias="scanType") + scan_type: ScanType concurrent: int = Field(ge=0) From 49af44149d57a5d08412f5150c24fa425d8ea99e Mon Sep 17 00:00:00 2001 From: Mark Feldhousen Date: Fri, 18 Oct 2024 12:23:11 -0400 Subject: [PATCH 132/139] Remove inexplicable quoting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 😜 Co-authored-by: dav3r --- src/cyhy_db/models/scan_doc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index ddf8510..feee961 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -100,7 +100,7 @@ async def reset_latest_flag_by_ip( @classmethod async def tag_latest( - cls, owners: List[str], snapshot: Union["SnapshotDoc", ObjectId, str] + cls, owners: List[str], snapshot: Union[SnapshotDoc, ObjectId, str] ): """Tag the latest scan for given owners with a snapshot id.""" from . import SnapshotDoc From a593d2de37d09a0e87fd141c9f2668a67be4b282 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 18 Oct 2024 17:58:27 -0400 Subject: [PATCH 133/139] Add asyncio_default_fixture_loop_scope configuration option to pytest.ini --- pytest.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/pytest.ini b/pytest.ini index ae07d53..505a21e 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,3 +1,4 @@ [pytest] addopts = -v -ra --cov --log-cli-level=INFO asyncio_mode = auto +asyncio_default_fixture_loop_scope = session From 86b493af4a22147e610edc045962cc0c0a120ae2 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 18 Oct 2024 18:02:20 -0400 Subject: [PATCH 134/139] Enable populate_by_name to handle aliases I'm not sure if this is a work-around for a bug, or a needed option in Pydantic. See: - https://github.com/BeanieODM/beanie/issues/369 - https://docs.pydantic.dev/latest/api/config/#pydantic.config.ConfigDict.populate_by_name --- src/cyhy_db/models/place_doc.py | 2 +- src/cyhy_db/models/ticket_doc.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/cyhy_db/models/place_doc.py b/src/cyhy_db/models/place_doc.py index a170b0a..ed1e356 100644 --- a/src/cyhy_db/models/place_doc.py +++ b/src/cyhy_db/models/place_doc.py @@ -11,7 +11,7 @@ class PlaceDoc(Document): """The place document model.""" - model_config = ConfigDict(extra="forbid") + model_config = ConfigDict(extra="forbid", populate_by_name=True) class_: str = Field(alias="class") # 'class' is a reserved keyword in Python country_name: str diff --git a/src/cyhy_db/models/ticket_doc.py b/src/cyhy_db/models/ticket_doc.py index a4d7deb..f621216 100644 --- a/src/cyhy_db/models/ticket_doc.py +++ b/src/cyhy_db/models/ticket_doc.py @@ -30,7 +30,9 @@ class EventDelta(BaseModel): """The event delta model.""" - from_: Optional[bool | float | int | str] = Field(..., serialization_alias="from") + model_config = ConfigDict(populate_by_name=True) + + from_: Optional[bool | float | int | str] = Field(..., alias="from") key: str = Field(...) to: Optional[bool | float | int | str] = Field(...) From d338efe50bcfb4d615282af3e5c60661cbec4bf3 Mon Sep 17 00:00:00 2001 From: Felddy Date: Fri, 18 Oct 2024 18:03:59 -0400 Subject: [PATCH 135/139] Add ScanLimit save test --- tests/test_request_doc.py | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/test_request_doc.py b/tests/test_request_doc.py index aa2f072..5837437 100644 --- a/tests/test_request_doc.py +++ b/tests/test_request_doc.py @@ -8,7 +8,8 @@ # cisagov Libraries from cyhy_db.models import RequestDoc -from cyhy_db.models.request_doc import Agency, Window +from cyhy_db.models.enum import ScanType +from cyhy_db.models.request_doc import Agency, ScanLimit, Window async def test_init(): @@ -60,3 +61,23 @@ def test_parse_time_invalid_type(): match="Invalid time format. Expected a string in '%H:%M:%S' format or datetime.time instance.", ): Window.parse_time(invalid_time_type) + + +async def test_scan_limit(): + """Test the ScanLimit model.""" + # Create a RequestDoc object + request_doc = RequestDoc( + agency=Agency(name="Office of Fragile Networking", acronym="OFN") + ) + + scan_limit = ScanLimit(scan_type=ScanType.CYHY, concurrent=1) + assert scan_limit.scan_type == ScanType.CYHY, "Scan type was not set correctly" + assert scan_limit.concurrent == 1, "Concurrent was not set correctly" + + request_doc.scan_limits.append(scan_limit) + assert ( + request_doc.scan_limits[0].scan_type == ScanType.CYHY + ), "Scan type was not set correctly" + await request_doc.save() + + # TODO complete this test From 84254b94ac6c5050e74dfe7056dd04faa0ff4687 Mon Sep 17 00:00:00 2001 From: Felddy Date: Sat, 19 Oct 2024 15:03:15 -0400 Subject: [PATCH 136/139] Attempt to improve ScanDoc test coverage In an attempt to reach 100% coverage I broke out the various ways `reset_latestest_by_ip` could be called... to no avail. I believe this issue explains why the list comprehension shows a branch part as uncovered: - https://github.com/nedbat/coveragepy/issues/1617 --- src/cyhy_db/models/scan_doc.py | 8 ++- tests/test_scan_doc.py | 106 ++++++++++++++++++++++++--------- 2 files changed, 83 insertions(+), 31 deletions(-) diff --git a/src/cyhy_db/models/scan_doc.py b/src/cyhy_db/models/scan_doc.py index feee961..a552714 100644 --- a/src/cyhy_db/models/scan_doc.py +++ b/src/cyhy_db/models/scan_doc.py @@ -87,10 +87,12 @@ async def reset_latest_flag_by_ip( ), ): """Reset the latest flag for all scans for a given IP address.""" - if isinstance(ips, Iterable): - ip_ints = [int(ip_address(x)) for x in ips] - else: + if isinstance(ips, (int, IPv4Address, str)): ip_ints = [int(ip_address(ips))] + else: + # Are you questing for 100% test coverage? + # Here be dragons: https://github.com/nedbat/coveragepy/issues/1617 + ip_ints = [int(ip_address(x)) for x in ips] # flake8 E712 is "comparison to True should be 'if cond is True:' or 'if # cond:'" but this is unavoidable due to Beanie syntax. diff --git a/tests/test_scan_doc.py b/tests/test_scan_doc.py index 79d80e6..4b69c97 100644 --- a/tests/test_scan_doc.py +++ b/tests/test_scan_doc.py @@ -141,34 +141,6 @@ async def test_reset_latest_flag_by_owner(): assert scan_doc.latest is False -async def test_reset_latest_flag_by_ip(): - """Test resetting the latest flag by IP address. - - This test verifies that the latest flag of ScanDoc objects is correctly - reset when the reset_latest_flag_by_ip method is called. - """ - IP_TO_RESET_1 = ipaddress.ip_address("128.205.1.2") - IP_TO_RESET_2 = ipaddress.ip_address("128.205.1.3") - scan_doc_1 = ScanDoc(ip=IP_TO_RESET_1, owner="RESET_BY_IP", source="nmap") - scan_doc_2 = ScanDoc(ip=IP_TO_RESET_2, owner="RESET_BY_IP", source="nmap") - await scan_doc_1.save() - await scan_doc_2.save() - # Check that the latest flag is set to True - assert scan_doc_1.latest is True - # Reset the latest flag on single IP - await ScanDoc.reset_latest_flag_by_ip(IP_TO_RESET_1) - # Retrieve the ScanDoc object from the database - await scan_doc_1.sync() - # Check that the latest flag is set to False - assert scan_doc_1.latest is False - # Reset by both IPs - await ScanDoc.reset_latest_flag_by_ip([IP_TO_RESET_1, IP_TO_RESET_2]) - # Retrieve the ScanDoc object from the database - await scan_doc_2.sync() - # Check that the latest flag is set to False - assert scan_doc_2.latest is False - - async def test_tag_latest_snapshot_doc(): """Test tagging the latest scan with a SnapshotDoc. @@ -286,3 +258,81 @@ async def test_tag_latest_invalid_type(): # Confirm that the scan does not have a snapshot assert scan_doc.snapshots == [], "Scan should not have any snapshots" + + +async def test_reset_latest_flag_by_ip_single(): + """Test reset_latest_flag_by_ip with a single IP address.""" + owner = "RESET_FLAG_SINGLE_IP" + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner=owner, + source="nmap", + ) + await scan_doc.save() + + # Reset the latest flag for a single IP address + await ScanDoc.reset_latest_flag_by_ip(scan_doc.ip) + + # Retrieve the ScanDoc object from the database + scan_doc = await ScanDoc.find_one(ScanDoc.id == scan_doc.id) + + # Check that the latest flag has been reset + assert ( + scan_doc.latest is False + ), "The latest flag was not reset for the single IP address" + + +async def test_reset_latest_flag_by_ip_list(): + """Test reset_latest_flag_by_ip with a list of IP addresses.""" + owner = "RESET_FLAG_IP_LIST" + scan_doc_1 = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner=owner, + source="nmap", + ) + await scan_doc_1.save() + + scan_doc_2 = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_2_STR), + owner=owner, + source="nmap", + ) + await scan_doc_2.save() + + # Reset the latest flag for a list of IP addresses + await ScanDoc.reset_latest_flag_by_ip([scan_doc_1.ip, scan_doc_2.ip]) + + # Retrieve the ScanDoc objects from the database + scan_doc_1 = await ScanDoc.find_one(ScanDoc.id == scan_doc_1.id) + scan_doc_2 = await ScanDoc.find_one(ScanDoc.id == scan_doc_2.id) + + # Check that the latest flag has been reset for both IP addresses + assert ( + scan_doc_1.latest is False + ), "The latest flag was not reset for the first IP address" + assert ( + scan_doc_2.latest is False + ), "The latest flag was not reset for the second IP address" + + +@pytest.mark.asyncio +async def test_reset_latest_flag_by_ip_empty_iterable(): + """Test reset_latest_flag_by_ip with an empty iterable.""" + owner = "RESET_FLAG_EMPTY_ITERABLE" + scan_doc = ScanDoc( + ip=ipaddress.ip_address(VALID_IP_1_STR), + owner=owner, + source="nmap", + ) + await scan_doc.save() + + # Reset the latest flag for an empty list of IP addresses + await ScanDoc.reset_latest_flag_by_ip([]) + + # Retrieve the ScanDoc object from the database + scan_doc = await ScanDoc.find_one(ScanDoc.id == scan_doc.id) + + # Check that the latest flag has not been modified + assert ( + scan_doc.latest is True + ), "The latest flag should remain True for empty iterable input" From d0106dce7fa5b673e499364b91c82d2dc2a9f69c Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 21 Oct 2024 10:22:53 -0400 Subject: [PATCH 137/139] Sort pytest options --- pytest.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytest.ini b/pytest.ini index 505a21e..caca126 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,4 +1,4 @@ [pytest] addopts = -v -ra --cov --log-cli-level=INFO -asyncio_mode = auto asyncio_default_fixture_loop_scope = session +asyncio_mode = auto From b1ffa0cbefee307b9f33c333412555156e4d19c7 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 21 Oct 2024 10:58:58 -0400 Subject: [PATCH 138/139] Update version file path in bump_version.sh --- bump_version.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bump_version.sh b/bump_version.sh index bd520bd..df7c371 100755 --- a/bump_version.sh +++ b/bump_version.sh @@ -6,7 +6,7 @@ set -o nounset set -o errexit set -o pipefail -VERSION_FILE=src/example/_version.py +VERSION_FILE=src/cyhy_db/_version.py HELP_INFORMATION="bump_version.sh (show|major|minor|patch|prerelease|build|finalize)" From 74a9505f336c3f235608570450e826b55d680f73 Mon Sep 17 00:00:00 2001 From: Felddy Date: Mon, 21 Oct 2024 10:59:19 -0400 Subject: [PATCH 139/139] Bump version from 0.0.1 to 1.0.0 --- src/cyhy_db/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cyhy_db/_version.py b/src/cyhy_db/_version.py index f102a9c..5becc17 100644 --- a/src/cyhy_db/_version.py +++ b/src/cyhy_db/_version.py @@ -1 +1 @@ -__version__ = "0.0.1" +__version__ = "1.0.0"