bug fixes

Author: toddbruner

src/app/api/__init__.py

@@ -3,13 +3,15 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.openapi.utils import get_openapi
+from fastapi.exception_handlers import http_exception_handler
 from fastapi.openapi.docs import (
     get_redoc_html,
     get_swagger_ui_html
 )
 from fastapi.staticfiles import StaticFiles
 from starlette.middleware.base import BaseHTTPMiddleware
 from starlette.responses import PlainTextResponse
+from starlette.exceptions import HTTPException
 from sqlalchemy.exc import TimeoutError, DatabaseError
 
 from app.api.endpoints import router
@@ -23,7 +25,10 @@ async def db_middleware(request: Request, call_next) -> None:
     """
     ROUTE_PREFIX_WHITELIST = ("/api/v1/firehose")
     if request.url.path.startswith(ROUTE_PREFIX_WHITELIST):
-        return await call_next(request)
+        try:
+            return await call_next(request)
+        except HTTPException as e:
+            return await http_exception_handler(request, e)
     else:
         try:
             with SessionLocal() as db_session:
@@ -41,7 +46,7 @@ async def db_middleware(request: Request, call_next) -> None:
         except TimeoutError:
             return PlainTextResponse("Timeout when connecting to database", status_code=503)
         except DatabaseError:
-            raise
+            return PlainTextResponse("Error when connecting to database", status_code=500)
 
     return response
 
@@ -51,31 +56,32 @@ def create_app() -> FastAPI:
     :return:
     """
     # Set up logging format
-    # dictConfig({
-    #    "version": 1,
-    #    "disable_existing_loggers": False,
-    #    "formatters": {
-    #        "access": {
-    #            "()": "uvicorn.logging.AccessFormatter",
-    #            "fmt": '%(asctime)s - %(levelprefix)s %(client_addr)s - \"%(request_line)s\" %(status_code)s',
-    #            "use_colors": True
-    #        },
-    #    },
-    #    "handlers": {
-    #        "access": {
-    #            "formatter": "access",
-    #            "class": "logging.StreamHandler",
-    #            "stream": "ext://sys.stdout",
-    #        },
-    #    },
-    #    "loggers": {
-    #        "uvicorn.access": {
-    #            "handlers": ["access"],
-    #            "level": "INFO",
-    #            "propagate": False
-    #        },
-    #    },
-    # })
+    dictConfig({
+       "version": 1,
+       "disable_existing_loggers": False,
+       "formatters": {
+           "access": {
+               "()": "uvicorn.logging.AccessFormatter",
+               "fmt": "%(levelprefix)s %(asctime)s.%(msecs)03d - %(client_addr)s - \"%(request_line)s\" %(status_code)s",
+               "datefmt": "%Y-%m-%dT%H:%M:%S",
+               "use_colors": True
+           },
+       },
+       "handlers": {
+           "access": {
+               "formatter": "access",
+               "class": "logging.StreamHandler",
+               "stream": "ext://sys.stdout",
+           },
+       },
+       "loggers": {
+           "uvicorn.access": {
+               "handlers": ["access"],
+               "level": "INFO",
+               "propagate": False
+           },
+       },
+    })
 
     app = FastAPI(
         debug=settings.DEBUG,
@@ -95,6 +101,7 @@ def create_app() -> FastAPI:
         allow_credentials=True,
         allow_methods=["*"],
         allow_headers=["*"],
+        expose_headers=["Content-Disposition"],  # allows for frontend to get download file names
     )
 
     app.add_middleware(

src/app/api/deps.py

@@ -174,7 +174,7 @@ def log(self, action, thing, thing_type=None, thing_pk=None, log_thing=True,
             # permissions (Role) - contains all permissions the role is a
             #       part of (only for cascade purposes)
             excluded_fields = ["schema_column", "entity_type", "tag_type",
-                               "permissions"]
+                               "permissions", "entity"]
             data = json.dumps(
                 jsonable_encoder(thing.as_dict(exclude_keys=excluded_fields))
             )

src/app/api/endpoints/audit.py

@@ -73,6 +73,7 @@ def search_audits(
     - Non-numeric/non-date fields can't use the range filters, for example `subject` or `description`. If range filters are provided the system will treat them as a list filter instead.
     - If none of the range or list filters work it will attempt to do a normal search
     - Datetimes are parsed using the [dateutil module](https://dateutil.readthedocs.io/en/stable/parser.html#dateutil.parser.parse)
+        - Some fields (e.g. event subjects and entity values) default to "contains" string searches, where any item containing the searched string will match (so searching for `example.com` would match both `example.com` and `foo.example.com`). Searching these fields with list searches (with square brackets) disables this feature for all list items, but "range" searches (with parentheses) search normally. For example, searching `[example.com]` wouldn't match `foo.example.com` but searching `(example.com)` would match.
     """
     try:
         filter_dict = get_search_filters(search_schema)

src/app/api/endpoints/file.py

@@ -122,11 +122,12 @@ def _download(db: Session, audit_logger: deps.AuditLogger, ids: list[int], passw
 
         # if only one file to download and no password was provided send the file
         if len(ids) == 1 and password is None:
+            filename = quote(fileobj.filename.encode("utf-8"))
             return StreamingResponse(
                 filestream,
                 media_type=fileobj.content_type,
                 headers={
-                    "content-disposition": f"attachment; filename*=UTF-8''{quote(fileobj.filename.encode('utf-8'))}"
+                    "content-disposition": f'attachment; filename="{filename}"'
                 },
             )
         # otherwise add file to zip

src/app/api/endpoints/firehose.py

@@ -14,18 +14,10 @@
 from app.schemas import TokenPayload
 
 router = APIRouter()
-
-# Close all connections on SIGINT, SIGTERM (more graceful shutdown)
+shutdown_signal_set = False
 shutdown = False
 
 
-def signal_shutdown(*args):
-    global shutdown
-    shutdown = True
-    for task in asyncio.all_tasks():
-        task.cancel()
-
-
 @router.get('/', summary="Stream audit events")
 async def stream_audits(
     *,
@@ -36,9 +28,24 @@ async def stream_audits(
     Stream a list of changes to SCOT data in real time (e.g. creation and
     modification events), as well as notifications for the user
     """
-    if signal.getsignal(signal.SIGINT) != signal_shutdown:
-        signal.signal(signal.SIGINT, signal_shutdown)
-        signal.signal(signal.SIGTERM, signal_shutdown)
+    # Close all connections on SIGINT, SIGTERM (more graceful shutdown)
+    def get_shutdown_signal(signal_type):
+        oldsignal = signal.getsignal(signal_type)
+        if oldsignal == signal.SIG_IGN or oldsignal == signal.SIG_DFL:
+            oldsignal = None
+
+        def signal_shutdown(*args):
+            global shutdown
+            shutdown = True
+            if oldsignal:
+                oldsignal(*args)
+        return signal_shutdown
+
+    global shutdown_signal_set
+    if not shutdown_signal_set:
+        signal.signal(signal.SIGINT, get_shutdown_signal(signal.SIGINT))
+        signal.signal(signal.SIGTERM, get_shutdown_signal(signal.SIGTERM))
+        shutdown_signal_set = True
 
     async def event_generator():
         audit_checkpoint = None
@@ -103,6 +110,6 @@ async def event_generator():
                         yield json.dumps({'what': 'create', 'when': record[1].strftime('%Y-%m-%d %H:%M:%s'), 'element_type': 'notification', 'element_id': record[0], 'username': current_user.username})
                 await asyncio.sleep(2)
         except asyncio.CancelledError as e:
-            pass
+            raise
 
     return EventSourceResponse(event_generator())

src/app/api/endpoints/generic.py

@@ -923,6 +923,7 @@ def search_object(
         - Non-numeric/non-date fields can't use the range filters, for example `subject` or `description`. If range filters are provided the system will treat them as a list filter instead.
         - If none of the range or list filters work it will attempt to do a normal search
         - Datetimes are parsed using the [dateutil module](https://dateutil.readthedocs.io/en/stable/parser.html#dateutil.parser.parse)
+        - Some fields (e.g. event subjects and entity values) default to "contains" string searches, where any item containing the searched string will match (so searching for `example.com` would match both `example.com` and `foo.example.com`). Searching these fields with list searches (with square brackets) disables this feature for all list items, but "range" searches (with parentheses) search normally. For example, searching `[example.com]` wouldn't match `foo.example.com` but searching `(example.com)` would match.
         """
 
         if target_type in deps.PermissionCheck.type_allow_whitelist:

src/app/crud/base.py

@@ -215,14 +215,20 @@ def _str_filter(self, query: Query, filter_dict: dict, column: str, escape: bool
             column_obj = filter_dict.pop(column, None)
             if column_obj is not None:
                 # does this even make sense? raise exception or just treat it like range?
-                if isinstance(column_obj, tuple) or isinstance(column_obj, list):
+                if isinstance(column_obj, tuple):
                     condition = []
                     for item in column_obj:
                         if escape:
                             condition.append(model.like(f"%{escape_sql_like(item)}%"))
                         else:
                             condition.append(model.like(f"%{item}%"))
                     query = query.filter(or_(*condition))
+                # If it's a list, disable string contains querying
+                elif isinstance(column_obj, list):
+                    condition = []
+                    for item in column_obj:
+                        condition.append(model == item)
+                    query = query.filter(or_(*condition))
                 else:
                     if escape:
                         t = escape_sql_like(column_obj)

src/app/crud/crud_alertgroup.py

@@ -151,111 +151,6 @@ def create(
             audit_logger.log("create", db_obj)
         return db_obj
 
-    def create_with_permissions(
-        self,
-        db_session: Session,
-        *,
-        obj_in: Union[AlertGroupCreate, AlertGroupDetailedCreate],
-        perm_in: dict[PermissionEnum, list],
-        audit_logger=None,
-    ) -> AlertGroup:
-        """Create an alertgroup given an AlertGroupCreate Object
-        First create the alertgroup. Get back the alertgroup ID.
-        Next, create the alertgroup schema
-        Next, create the alerts, with the given alertgroup ID.
-        Fin
-
-        """
-
-        if PermissionEnum.admin in perm_in:
-            raise ValueError("Users cannot assign admin permissions")
-        db_obj = AlertGroup(
-            owner=obj_in.owner,
-            tlp=obj_in.tlp,
-            view_count=obj_in.view_count,
-            first_view=obj_in.first_view,
-            message_id=obj_in.message_id,
-            subject=obj_in.subject,
-        )
-
-        db_session.add(db_obj)
-        db_session.flush()
-        db_session.refresh(db_obj)
-
-        tt = self.model.target_type_enum()
-        # Assign permissions (if applicable)
-        if tt:
-            # need to import here to avoid circular dependency
-            from app.crud import permission, role
-
-            for perm in perm_in:
-                new_perm = {
-                    "permission": perm,
-                    "target_type": tt,
-                    "target_id": db_obj.id,
-                }
-                for r in perm_in[perm]:
-                    role_id = r
-                    if not isinstance(r, int):
-                        role_id = role.get_role_by_name(db_session, r).id
-                    new_perm["role_id"] = role_id
-                    permission.create(
-                        db_session, obj_in=new_perm, audit_logger=audit_logger
-                    )
-
-        # Also create the schema and add alerts if detailed create
-        if isinstance(obj_in, AlertGroupDetailedCreate):
-            schema = None
-            if obj_in.alert_schema:
-                schema = obj_in.alert_schema
-            # Create the schema if one wasn't given
-            elif obj_in.alerts:
-                schema = self.validate_alerts(obj_in.alerts)
-            if schema is not None:
-                # We have a valid schema, now let's add it
-                for schema_column in schema:
-                    schema_column.alertgroup_id = db_obj.id
-                    crud.alert_group_schema.create(
-                        db_session=db_session, obj_in=schema_column
-                    )
-                # We've now validated and added the schema, lets add the alerts now
-                for alert in obj_in.alerts:
-                    if alert.owner is None:
-                        alert.owner = obj_in.owner
-                    if alert.tlp == TlpEnum.unset:
-                        alert.tlp = obj_in.tlp
-                    alert.alertgroup_id = db_obj.id
-                    crud.alert.create_with_permissions(db_session=db_session, obj_in=alert, perm_in=perm_in)
-        self.publish("create", db_obj)
-        # Last thing, check to see if the alert group subject is contained in any existing signatures, if so, add links.
-        looking_for = re.sub(r"\([^()]*\)", "", obj_in.subject)
-        if "*" in looking_for or "_" in looking_for:
-            looking_for = looking_for.lstrip().replace("\\", "\\\\")
-            looking_for = (
-                looking_for.replace("_", "__").replace("*", "%").replace("?", "_")
-            )
-        else:
-            looking_for = looking_for.lstrip()
-            looking_for = "%{0}%".format(escape_sql_like(looking_for))
-        sigs_to_link_query = db_session.query(Signature).filter(
-            Signature.name.ilike(looking_for)
-        )
-        results = sigs_to_link_query.all()
-        for sig_to_link in results:
-            crud.link.create(
-                db_session=db_session,
-                obj_in=LinkCreate(
-                    v0_type=TargetTypeEnum.alertgroup,
-                    v0_id=db_obj.id,
-                    v1_type=TargetTypeEnum.signature,
-                    v1_id=sig_to_link.id,
-                    context=f"Automatically linked from alertgroup subject matching: {looking_for}",
-                ),
-            )
-        if audit_logger is not None:
-            audit_logger.log("create", db_obj)
-        return db_obj
-
     def create_with_owner(
         self,
         db_session: Session,

src/app/crud/crud_entity.py

@@ -384,7 +384,7 @@ def add_enrichment(
         db_session.refresh(entity)
         db_session.flush()
         if audit_logger is not None:
-            audit_logger.log("create", enrichment)
+            audit_logger.log("create", enrichment, log_thing=False)
         return entity
 
     def add_entity_classes(

src/app/db/base_class.py

@@ -20,7 +20,7 @@ class Base:
     def __tablename__(cls) -> str:
         return cls.__name__.lower()
 
-    def as_dict(self, exclude_keys=["entity_type", "tag_type", "entity_types", "entity_classes", "permissions"], pretty_keys: bool = False, enum_value: bool = False):
+    def as_dict(self, exclude_keys=["entity_type", "tag_type", "entity_types", "entity_classes", "permissions", "entity"], pretty_keys: bool = False, enum_value: bool = False):
         """
         Serializes this model to a dictionary as accurately as possible
         Contains logic to prevent infinite recursion for circular references
@@ -132,4 +132,4 @@ def get_model_by_target_type(cls, target_type: TargetTypeEnum):
                 registry_instance = getattr(cls, "registry")
                 for mapper_ in registry_instance.mappers:
                     if (mapper_.class_.__tablename__ == table_name):
-                        return mapper_.class_
+                        return mapper_.class_