Merge pull request #138 from blacklanternsecurity/dev

Dev -> Stable

Author: TheTechromancer

bbot_server/modules/__init__.py

@@ -10,7 +10,7 @@
 
 # needed for asset model preloading
 from bbot_server.assets import CustomAssetFields  # noqa: F401
-from typing import List, Optional, Dict, Any, Annotated  # noqa: F401
+from typing import List, Literal, Optional, Dict, Any, Annotated  # noqa: F401
 from pydantic import Field, BeforeValidator, AfterValidator, UUID4  # noqa: F401
 
 log = logging.getLogger(__name__)

bbot_server/modules/activity/activity_models.py

@@ -43,7 +43,7 @@ class Activity(BaseHostModel):
         default_factory=utc_now, description="Time when this activity was created"
     )
     archived: Annotated[bool, "indexed"] = False
-    description: Annotated[str, "indexed"]
+    description: Annotated[str, "indexed", "indexed-text"]
     description_colored: str = Field(default="")
     detail: dict[str, Any] = {}
     module: Annotated[Optional[str], "indexed"] = None

bbot_server/modules/assets/assets_models.py

@@ -14,9 +14,6 @@ class AdvancedAssetQuery(AssetQuery):
 
     async def build(self, applet=None):
         query = await super().build(applet)
-        print(f"QUERERY BEFOREE: {query}")
-        print(f"SELF>TTTYPE: {self.type}")
         if ("type" not in query) and self.type:
             query["type"] = self.type
-        print(f"UQYWERQWEYRYRY: {query}")
         return query

bbot_server/modules/findings/findings_api.py

@@ -5,13 +5,31 @@
 from bbot_server.applets.base import BaseApplet, api_endpoint
 from bbot_server.modules.findings.findings_models import Finding, SEVERITY_COLORS, SeverityScore, FindingsQuery
 
+# Max CVSS score for each severity band (top of range).
+# Used to derive a default risk score from finding_max_severity.
+SEVERITY_TO_CVSS = {
+    "INFO": 0.0,
+    "LOW": 0.1,
+    "MEDIUM": 4.0,
+    "HIGH": 7.0,
+    "CRITICAL": 9.0,
+}
+
 
 # add 'findings' field to the main asset model
 class FindingFields(CustomAssetFields):
     findings: Annotated[list[str], "indexed", "indexed-text"] = []
     finding_severities: Annotated[dict[str, int], "indexed"] = {}
-    finding_max_severity: Optional[Annotated[str, "indexed"]] = None
+    finding_max_severity: Annotated[Optional[str], "indexed"] = None
     finding_max_severity_score: Annotated[int, "indexed"] = 0
+    # Effective risk score for this asset: None or a float from 0.0 to 10.0
+    # (1 decimal place). Auto-synced from finding_max_severity (using CVSS
+    # thresholds) unless risk_override is True.
+    risk: Annotated[Optional[float], "indexed"] = None
+    # Whether risk has been manually overridden. When True, new findings
+    # will NOT auto-update risk. Clearing the override resets this to False
+    # and reverts risk to the CVSS-derived value.
+    risk_override: Annotated[bool, "indexed"] = False
 
 
 class FindingsApplet(BaseApplet):
@@ -134,6 +152,68 @@ async def severity_counts(
         findings = dict(sorted(findings.items(), key=lambda x: x[1], reverse=True))
         return findings
 
+    @api_endpoint("/set_risk", methods=["PATCH"], summary="Set or clear a manual risk score for an asset")
+    async def set_risk(
+        self,
+        host: Annotated[str, Query(description="The host of the asset to update")],
+        risk: Annotated[
+            Optional[float],
+            Query(
+                description=(
+                    "Risk score from 0.0 to 10.0 (1 decimal place). "
+                    "Omit to clear the override and revert to the auto-calculated CVSS value."
+                )
+            ),
+        ] = None,
+        override_none: Annotated[
+            bool,
+            Query(
+                description=(
+                    "Set to true to explicitly override risk to None (no risk score). "
+                    "Takes precedence over the risk parameter."
+                )
+            ),
+        ] = False,
+    ) -> dict:
+        """
+        Manually set or clear an asset's risk score.
+
+        Three modes:
+          - risk=<float>       → override risk to the given value (0.0–10.0, 1 decimal).
+          - override_none=true → override risk to None (e.g. "no risk score").
+          - (omit both)        → clear the override and revert to the CVSS-derived
+                                 value from finding_max_severity.
+        """
+        asset = await self.root._get_asset(host=host, fields=["finding_max_severity"])
+        if not asset:
+            raise self.BBOTServerNotFoundError(f"Asset {host} not found")
+
+        if override_none:
+            # Explicit override to None
+            update = {"risk": None, "risk_override": True}
+            description = f"Risk manually set to [bold]None[/bold] on [bold]{host}[/bold]"
+        elif risk is not None:
+            # Override to a specific float value
+            if risk < 0.0 or risk > 10.0:
+                raise self.BBOTServerValueError("risk must be between 0.0 and 10.0")
+            risk = round(risk, 1)
+            update = {"risk": risk, "risk_override": True}
+            description = f"Risk manually set to [bold]{risk}[/bold] on [bold]{host}[/bold]"
+        else:
+            # Clear the override: revert to CVSS-derived value
+            finding_max_severity = asset.get("finding_max_severity", None)
+            reverted_risk = SEVERITY_TO_CVSS.get(finding_max_severity) if finding_max_severity else None
+            update = {"risk": reverted_risk, "risk_override": False}
+            description = f"Risk override cleared on [bold]{host}[/bold], reverted to [bold]{reverted_risk}[/bold]"
+
+        await self.root._update_asset(host, update)
+        await self.emit_activity(
+            type="RISK_UPDATED",
+            description=description,
+            detail={"host": host, **update},
+        )
+        return {"host": host, "risk": update["risk"], "risk_override": update["risk_override"]}
+
     async def handle_event(self, event, asset):
         name = event.data_json["name"]
         description = event.data_json["description"]
@@ -164,8 +244,7 @@ async def compute_stats(self, asset, stats):
             - finding names
             - finding severities
             - finding hosts
-            - finding max severity
-            - finding max severity score
+            - severity counts by host
         """
         finding_names = getattr(asset, "findings", [])
         finding_severities = getattr(asset, "finding_severities", {})
@@ -181,24 +260,13 @@ async def compute_stats(self, asset, stats):
         for finding_severity, count in finding_severities.items():
             severity_stats[finding_severity] = severity_stats.get(finding_severity, 0) + count
 
-        max_severity_score = max([asset.finding_max_severity_score, finding_stats.get("max_severity_score", 0)])
-        finding_stats["max_severity_score"] = max_severity_score
-        if max_severity_score > 0:
-            max_severity = SeverityScore.to_str(max_severity_score)
-        else:
-            max_severity = None
-        finding_stats["max_severity"] = max_severity
-
-        if asset.finding_max_severity_score > 0:
-            severities_by_host[asset.host] = {
-                "max_severity": asset.finding_max_severity,
-                "max_severity_score": asset.finding_max_severity_score,
-            }
+        if finding_severities:
+            severities_by_host[asset.host] = dict(sorted(finding_severities.items(), key=lambda x: x[1], reverse=True))
 
         finding_stats["names"] = dict(sorted(name_stats.items(), key=lambda x: x[1], reverse=True))
         finding_stats["counts_by_host"] = dict(sorted(counts_by_host.items(), key=lambda x: x[1], reverse=True))
         finding_stats["severities_by_host"] = dict(
-            sorted(severities_by_host.items(), key=lambda x: x[1]["max_severity_score"], reverse=True)
+            sorted(severities_by_host.items(), key=lambda x: sum(x[1].values()), reverse=True)
         )
         finding_stats["severities"] = dict(sorted(severity_stats.items(), key=lambda x: x[1], reverse=True))
         stats["findings"] = finding_stats
@@ -245,17 +313,36 @@ async def _insert_or_update_finding(self, finding: Finding, asset, event=None):
         else:
             asset.finding_max_severity_score = 0
             asset.finding_max_severity = None
+        # Auto-sync risk from finding_max_severity when not manually overridden.
+        old_risk = getattr(asset, "risk", None)
+        if not getattr(asset, "risk_override", False):
+            if asset.finding_max_severity is not None:
+                asset.risk = SEVERITY_TO_CVSS[asset.finding_max_severity]
+            else:
+                asset.risk = None
 
         # insert the new vulnerability
         await self.root._insert_asset(finding.model_dump())
 
         severity_color = SEVERITY_COLORS[finding.severity_score]
 
-        return [
+        activities = [
             self.make_activity(
                 type="NEW_FINDING",
                 description=f"New finding with severity [bold {severity_color}]{finding.severity}[/bold {severity_color}]: [[bold {severity_color}]{finding.name}[/bold {severity_color}]] on [bold]{finding.host}[/bold]",
                 event=event,
                 detail=finding.model_dump(),
             )
         ]
+
+        # emit RISK_UPDATED if risk actually changed
+        if asset.risk != old_risk:
+            activities.append(
+                self.make_activity(
+                    type="RISK_UPDATED",
+                    description=f"Risk updated from [bold]{old_risk}[/bold] to [bold]{asset.risk}[/bold] on [bold]{asset.host}[/bold]",
+                    detail={"host": asset.host, "risk": asset.risk, "old_risk": old_risk},
+                )
+            )
+
+        return activities

bbot_server/modules/server/server_cli.py

@@ -6,6 +6,8 @@
 from subprocess import run
 from contextlib import suppress
 
+from pymongo import uri_parser
+
 from bbot_server.config import BBOT_SERVER_CONFIG as bbcfg, BBOT_SERVER_DIR
 from bbot_server.cli.base import BaseBBCTL, subcommand, Option, Annotated
 
@@ -158,7 +160,7 @@ def cleardb(
             )
 
         for store_name, store_config, data_desc in stores_to_clear:
-            db_name = store_config.uri.split("/")[-1]
+            db_name = uri_parser.parse_uri(store_config.uri)["database"]
             prefix = store_config.collection_prefix
             if not db_name:
                 raise self.BBOTServerError(f"{store_name.title()} database not found in config")

bbot_server/modules/targets/targets_api.py

@@ -113,7 +113,7 @@ async def refresh_asset_scope(self, host: str, target: BBOTTarget, target_id: UU
         scope_result_type = getattr(scope_result, "type", None)
         if scope_result_type == "NEW_IN_SCOPE_ASSET":
             asset_scope = sorted(set(asset_scope) | set([target_id]))
-        else:
+        elif scope_result_type == "ASSET_SCOPE_CHANGED":
             asset_scope = sorted(set(asset_scope) - set([target_id]))
         asset_results = await self.root.assets.collection.update_many(
             {"host": host},
@@ -358,17 +358,17 @@ async def _check_scope(self, host, resolved_hosts, target: BBOTTarget, target_id
         try:
             # we take the main host and its A/AAAA DNS records into account
             for rdtype, hosts in resolved_hosts.items():
-                for host in hosts:
+                for h in hosts:
                     # if any of the hosts are blacklisted, abort immediately
-                    if target.blacklisted(host):
-                        blacklisted_reason = f"{rdtype}->{host}"
+                    if target.blacklisted(h):
+                        blacklisted_reason = f"{rdtype}->{h}"
                         in_target_reason = ""
                         # break out of the loop
                         raise BlacklistedError
                     # check against whitelist
                     if not in_target_reason:
-                        if target.in_target(host):
-                            in_target_reason = f"{rdtype}->{host}"
+                        if target.in_target(h):
+                            in_target_reason = f"{rdtype}->{h}"
         except BlacklistedError:
             pass
 

bbot_server/store.py

@@ -7,7 +7,7 @@
 class BaseMongoStore(BaseDB):
     async def setup(self):
         self.client = AsyncMongoClient(self.uri)
-        self.db = self.client.get_database(self.db_name)
+        self.db = self.client.get_default_database()
         self.collection_prefix = getattr(self.db_config, "collection_prefix", "")
         bucket_name = f"{self.collection_prefix}fs" if self.collection_prefix else "fs"
         self.fs = AsyncGridFSBucket(self.db, bucket_name=bucket_name)

bbot_server/utils/misc.py

@@ -255,12 +255,13 @@ def combine_pydantic_models(models, model_name, base_model=BaseModel):
 
 
 # fmt: off
+# 20260417: removed $jsonSchema because it may reveal internal/private fields
 ALLOWED_QUERY_OPERATORS = {
     # Query Operators (excluding $where, $expr)
     "$eq", "$gt", "$gte", "$in", "$lt", "$lte", "$ne", "$nin",
     "$and", "$not", "$nor", "$or",
     "$exists", "$type",
-    "$jsonSchema", "$mod", "$search", "$text", "$regex",
+    "$mod", "$search", "$text", "$regex",
     "$geoIntersects", "$geoWithin", "$near", "$nearSphere",
     "$all", "$elemMatch", "$size",
     "$bitsAllClear", "$bitsAllSet", "$bitsAnyClear", "$bitsAnySet",
@@ -289,6 +290,7 @@ def _sanitize_mongo_query(data: Any) -> Any:
 
 
 # fmt: off
+# 20260417: removed $unionWith because it may allow fetches to other collections
 ALLOWED_AGG_OPERATORS = {
     # We intentionally exclude $match because it"s automatically added and sanitized separately
 
@@ -299,7 +301,7 @@ def _sanitize_mongo_query(data: Any) -> Any:
     "$planCacheStats", "$project", "$redact",
     "$replaceRoot", "$replaceWith", "$sample", "$search", "$searchMeta",
     "$set", "$setWindowFields", "$skip", "$sort", "$sortByCount",
-    "$unionWith", "$unset", "$unwind",
+    "$unset", "$unwind",
 
     # Aggregation Expression Operators (excluding $function, $accumulator)
     # Arithmetic

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bbot-server"
-version = "0.1.3"
+version = "0.1.4"
 description = ""
 authors = [{name = "TheTechromancer"}]
 license = "AGPL-3.0"

tests/test_applets/test_applet_assets.py

@@ -346,3 +346,83 @@ async def test_applet_target_filter(bbot_server, bbot_events):
     assert set(assets) == all_hosts_target2
     hosts = await bbot_server.get_hosts(target_id=target.id)
     assert set(hosts) == all_hosts_target2
+
+
+# test to make sure custom attributes on assets are queryable
+async def test_applet_custom_attributes(bbot_server, bbot_events):
+    bbot_server = await bbot_server(needs_worker=True)
+
+    # skip testing of the http interface (since insertion isn't supported)
+    if not bbot_server.is_native:
+        return
+
+    # ingest BBOT events to create some assets
+    scan1_events, scan2_events = bbot_events
+    for e in scan1_events:
+        await bbot_server.insert_event(e)
+
+    # wait for events to be processed
+    await asyncio.sleep(INGEST_PROCESSING_DELAY)
+
+    # verify assets exist
+    hosts = set(await bbot_server.get_hosts())
+    assert "evilcorp.com" in hosts
+    assert "www.evilcorp.com" in hosts
+    assert "api.evilcorp.com" in hosts
+
+    # manually add custom attributes to some assets via the collection
+    collection = bbot_server.assets.collection
+    await collection.update_one(
+        {"host": "evilcorp.com", "type": "Asset"},
+        {"$set": {"custom_tag": "important", "risk_score": 95}},
+    )
+    await collection.update_one(
+        {"host": "www.evilcorp.com", "type": "Asset"},
+        {"$set": {"custom_tag": "important", "risk_score": 50}},
+    )
+    await collection.update_one(
+        {"host": "api.evilcorp.com", "type": "Asset"},
+        {"$set": {"custom_tag": "low-priority", "risk_score": 10}},
+    )
+
+    # query by custom attribute - exact match
+    results = [a async for a in bbot_server.query_assets(query={"custom_tag": "important"})]
+    assert {a["host"] for a in results} == {"evilcorp.com", "www.evilcorp.com"}
+
+    # query by custom attribute - comparison operator
+    results = [a async for a in bbot_server.query_assets(query={"risk_score": {"$gte": 50}})]
+    assert {a["host"] for a in results} == {"evilcorp.com", "www.evilcorp.com"}
+
+    # query combining custom attribute with built-in filters
+    results = [a async for a in bbot_server.query_assets(query={"custom_tag": "important", "host": "evilcorp.com"})]
+    assert len(results) == 1
+    assert results[0]["host"] == "evilcorp.com"
+
+    # verify custom fields are returned in query results
+    results = [a async for a in bbot_server.query_assets(query={"host": "api.evilcorp.com"})]
+    assert len(results) == 1
+    assert results[0]["custom_tag"] == "low-priority"
+    assert results[0]["risk_score"] == 10
+
+    # query for a custom attribute value that doesn't exist
+    results = [a async for a in bbot_server.query_assets(query={"custom_tag": "nonexistent"})]
+    assert results == []
+
+    # aggregation on custom attributes
+    agg_results = [
+        a
+        async for a in bbot_server.query_assets(
+            aggregate=[
+                {"$group": {"_id": "$custom_tag", "avg_risk": {"$avg": "$risk_score"}}},
+                {"$sort": {"_id": 1}},
+            ],
+        )
+    ]
+    assert len(agg_results) == 3
+    assert agg_results[0] == {"_id": None, "avg_risk": None}
+    assert agg_results[1] == {"_id": "important", "avg_risk": 72.5}
+    assert agg_results[2] == {"_id": "low-priority", "avg_risk": 10.0}
+
+    # count with custom attribute filter
+    count = await bbot_server.count_assets(query={"custom_tag": "important"})
+    assert count == 2