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