From e20b4bb4f871840f0a68dbe2d710fbd281051ed0 Mon Sep 17 00:00:00 2001
From: Robin Schulz <hello@rschu.me>
Date: Sat, 30 Mar 2024 17:40:15 +0100
Subject: [PATCH] - missing typings added - `httpurl_to_str` utils function
 removed - `CustomBaseModel` created and used for all Pydantic model classes -
 `RandomAuditLogSettings` model renamed to `BulkAuditLogOptions`

---
 audit_logger/exceptions.py             |  2 +-
 audit_logger/main.py                   |  8 ++++----
 audit_logger/models/__init__.py        |  2 +-
 audit_logger/models/actor_details.py   |  6 ++++--
 audit_logger/models/audit_log_entry.py |  9 +++------
 audit_logger/models/config.py          | 10 ++++++----
 audit_logger/models/custom_base.py     | 12 ++++++++++++
 audit_logger/models/env_vars.py        |  6 ++++--
 audit_logger/models/request.py         | 11 +++++------
 audit_logger/models/resource.py        |  6 ++++--
 audit_logger/models/response_models.py |  9 +++++----
 audit_logger/models/search_params.py   | 19 ++++++++-----------
 audit_logger/models/server_details.py  |  6 ++++--
 audit_logger/utils.py                  | 22 ++++++++++------------
 14 files changed, 71 insertions(+), 57 deletions(-)
 create mode 100644 audit_logger/models/custom_base.py

diff --git a/audit_logger/exceptions.py b/audit_logger/exceptions.py
index 809243a..d1de847 100644
--- a/audit_logger/exceptions.py
+++ b/audit_logger/exceptions.py
@@ -38,7 +38,7 @@ async def validation_exception_handler(
 
 
 class BulkLimitExceededError(HTTPException):
-    def __init__(self, limit: int):
+    def __init__(self, limit: int) -> None:
         super().__init__(
             status_code=400,
             detail=f"The maximum number of items allowed in a bulk operation is {limit}.",
diff --git a/audit_logger/main.py b/audit_logger/main.py
index 8a82fe1..beababd 100644
--- a/audit_logger/main.py
+++ b/audit_logger/main.py
@@ -14,8 +14,8 @@
 from audit_logger.middlewares import add_middleware
 from audit_logger.models import (
     AuditLogEntry,
+    BulkAuditLogOptions,
     GenericResponse,
-    RandomAuditLogSettings,
     SearchParamsV2,
     SearchResults,
 )
@@ -109,7 +109,7 @@ async def create_bulk_audit_log_entries(
     Raises:
         Union[HTTPException, BulkLimitExceededError]
     """
-    bulk_limit = 250
+    bulk_limit = 350
     if len(audit_logs) > bulk_limit:
         raise BulkLimitExceededError(limit=bulk_limit)
     return await process_audit_logs(
@@ -121,7 +121,7 @@ async def create_bulk_audit_log_entries(
 
 @app.post("/create/create-bulk-auto")
 async def create_fake_audit_log_entries(
-    settings: RandomAuditLogSettings,
+    options: BulkAuditLogOptions,
 ) -> GenericResponse:
     """
     Generates and stores a single random audit log entry.
@@ -135,7 +135,7 @@ async def create_fake_audit_log_entries(
     return await process_audit_logs(
         elastic,
         cast(str, env_vars.elastic_index_name),
-        generate_audit_log_entries_with_fake_data(settings),
+        generate_audit_log_entries_with_fake_data(options),
     )
 
 
diff --git a/audit_logger/models/__init__.py b/audit_logger/models/__init__.py
index 1fafc5a..5df1195 100644
--- a/audit_logger/models/__init__.py
+++ b/audit_logger/models/__init__.py
@@ -1,7 +1,7 @@
 from .actor_details import ActorDetails
 from .audit_log_entry import AuditLogEntry
 from .config import APIMiddlewares, AppConfig, CORSSettings
-from .request import RandomAuditLogSettings
+from .request import BulkAuditLogOptions
 from .resource import ResourceDetails
 from .response_models import (
     ActorDetails,
diff --git a/audit_logger/models/actor_details.py b/audit_logger/models/actor_details.py
index 0a4e714..d62c2bd 100644
--- a/audit_logger/models/actor_details.py
+++ b/audit_logger/models/actor_details.py
@@ -1,7 +1,9 @@
 from enum import Enum
 from typing import Optional
 
-from pydantic import BaseModel, Field, IPvAnyAddress
+from pydantic import Field, IPvAnyAddress
+
+from audit_logger.models.custom_base import CustomBaseModel
 
 
 class ActorType(str, Enum):
@@ -9,7 +11,7 @@ class ActorType(str, Enum):
     SYSTEM = "system"
 
 
-class ActorDetails(BaseModel):
+class ActorDetails(CustomBaseModel):
     identifier: Optional[str] = Field(
         default=None,
         description="Unique identifier of the actor. "
diff --git a/audit_logger/models/audit_log_entry.py b/audit_logger/models/audit_log_entry.py
index a4c5b04..ab617eb 100644
--- a/audit_logger/models/audit_log_entry.py
+++ b/audit_logger/models/audit_log_entry.py
@@ -1,13 +1,14 @@
 from typing import Any, Dict, Optional
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import Field
 
 from audit_logger.models.actor_details import ActorDetails
+from audit_logger.models.custom_base import CustomBaseModel
 from audit_logger.models.resource import ResourceDetails
 from audit_logger.models.server_details import ServerDetails
 
 
-class AuditLogEntry(BaseModel):
+class AuditLogEntry(CustomBaseModel):
     timestamp: Optional[str] = Field(
         default=None,
         description="The date and time when the event occurred, in ISO 8601 format.",
@@ -44,7 +45,3 @@ class AuditLogEntry(BaseModel):
     meta: Dict[str, Any] = Field(
         default={}, description="Optional metadata about the event."
     )
-
-    # Forbid extra fields and raise an exception if any are found.
-    class Config:
-        extra = Extra.forbid
diff --git a/audit_logger/models/config.py b/audit_logger/models/config.py
index ad11bed..6087e59 100644
--- a/audit_logger/models/config.py
+++ b/audit_logger/models/config.py
@@ -1,9 +1,11 @@
 from typing import List, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
+from audit_logger.models.custom_base import CustomBaseModel
 
-class CORSSettings(BaseModel):
+
+class CORSSettings(CustomBaseModel):
     allow_origins: Optional[List[str]] = Field(
         default=["*"],
         description="List of allowed origin URLs. Use '*' to allow all origins.",
@@ -22,11 +24,11 @@ class CORSSettings(BaseModel):
     )
 
 
-class APIMiddlewares(BaseModel):
+class APIMiddlewares(CustomBaseModel):
     cors: CORSSettings = Field(description="CORS middleware settings")
 
 
-class AppConfig(BaseModel):
+class AppConfig(CustomBaseModel):
     middlewares: Optional[APIMiddlewares] = Field(
         description="API Middlewares settings",
     )
diff --git a/audit_logger/models/custom_base.py b/audit_logger/models/custom_base.py
new file mode 100644
index 0000000..9a972cc
--- /dev/null
+++ b/audit_logger/models/custom_base.py
@@ -0,0 +1,12 @@
+from typing import Any
+
+from pydantic import BaseModel, Extra
+
+
+class CustomBaseModel(BaseModel):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+
+    # Forbid extra fields and raise an exception if any are found.
+    class Config:
+        extra = Extra.forbid
diff --git a/audit_logger/models/env_vars.py b/audit_logger/models/env_vars.py
index f3db090..972bd22 100644
--- a/audit_logger/models/env_vars.py
+++ b/audit_logger/models/env_vars.py
@@ -1,9 +1,11 @@
 from typing import Optional
 
-from pydantic import BaseModel, Field, HttpUrl, field_validator
+from pydantic import Field, HttpUrl, field_validator
 
+from audit_logger.models.custom_base import CustomBaseModel
 
-class EnvVars(BaseModel):
+
+class EnvVars(CustomBaseModel):
     elastic_index_name: str = Field(...)
     elastic_url: str = Field(...)
     elastic_username: Optional[str] = Field(None)
diff --git a/audit_logger/models/request.py b/audit_logger/models/request.py
index 6b3e40d..282f407 100644
--- a/audit_logger/models/request.py
+++ b/audit_logger/models/request.py
@@ -1,16 +1,15 @@
 from typing import Optional
 
-from pydantic import BaseModel, Extra, Field
+from pydantic import Field
 
+from audit_logger.models.custom_base import CustomBaseModel
 
-class RandomAuditLogSettings(BaseModel):
+
+class BulkAuditLogOptions(CustomBaseModel):
     bulk_count: Optional[int] = Field(
         1,
         ge=1,
         le=350,
         description="Specifies the number of fake audit log entries to generate.",
+        alias="bulk_limit",
     )
-
-    # Forbid extra fields and raise an exception if any are found.
-    class Config:
-        extra = Extra.forbid
diff --git a/audit_logger/models/resource.py b/audit_logger/models/resource.py
index 7e289cb..9d87a04 100644
--- a/audit_logger/models/resource.py
+++ b/audit_logger/models/resource.py
@@ -1,9 +1,11 @@
 from typing import Optional
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
+from audit_logger.models.custom_base import CustomBaseModel
 
-class ResourceDetails(BaseModel):
+
+class ResourceDetails(CustomBaseModel):
     type: Optional[str] = Field(
         default=None,
         description="Type of the resource that was acted upon.",
diff --git a/audit_logger/models/response_models.py b/audit_logger/models/response_models.py
index 433d20a..b16e456 100644
--- a/audit_logger/models/response_models.py
+++ b/audit_logger/models/response_models.py
@@ -1,12 +1,13 @@
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel, Field
+from pydantic import Field
 
 from audit_logger.models.actor_details import ActorDetails
+from audit_logger.models.custom_base import CustomBaseModel
 from audit_logger.models.resource import ResourceDetails
 
 
-class LogEntryDetails(BaseModel):
+class LogEntryDetails(CustomBaseModel):
     timestamp: Optional[str] = Field(
         default=None, description="The timestamp of the log entry."
     )
@@ -44,7 +45,7 @@ class LogEntryDetails(BaseModel):
     )
 
 
-class SearchResults(BaseModel):
+class SearchResults(CustomBaseModel):
     hits: int = Field(default=0, description="The number of hits for the search query.")
     docs: List[Any] = Field(
         default=[], description="A list of documents that match the search query."
@@ -54,7 +55,7 @@ class SearchResults(BaseModel):
     )
 
 
-class GenericResponse(BaseModel):
+class GenericResponse(CustomBaseModel):
     status: str = Field(default=None, description="The status of the response.")
     success_count: int = Field(
         default=None, description="The number of successful items in the response."
diff --git a/audit_logger/models/search_params.py b/audit_logger/models/search_params.py
index b767ecf..0c09517 100644
--- a/audit_logger/models/search_params.py
+++ b/audit_logger/models/search_params.py
@@ -3,14 +3,15 @@
 from enum import Enum
 from typing import Any, Dict, List, Optional, Union
 
-from pydantic import BaseModel, Extra, Field, field_validator, model_validator
+from pydantic import Field, field_validator, model_validator
 
+from audit_logger.models.custom_base import CustomBaseModel
 from audit_logger.utils import is_valid_ip_v4_address, validate_date
 
 logger = logging.getLogger("audit_logger")
 
 
-class MaxResultsMixin(BaseModel):
+class MaxResultsMixin(CustomBaseModel):
     max_results: Optional[int] = Field(
         500,
         ge=1,
@@ -70,7 +71,7 @@ class AggregationTypeEnum(str, Enum):
     DATE_HISTOGRAM = "date_histogram"
 
 
-class SearchFilterParams(BaseModel):
+class SearchFilterParams(CustomBaseModel):
     field: FieldIdentifierEnum = Field(
         default=None,
         description="",
@@ -135,16 +136,16 @@ def check_type_is_valid(cls, v):
         return v
 
 
-class AggregationFilterParams(BaseModel):
+class AggregationFilterParams(CustomBaseModel):
     range: Optional[Dict[str, Dict[str, str]]]
 
 
-class SubAggregationConfig(BaseModel):
+class SubAggregationConfig(CustomBaseModel):
     type: AggregationTypeEnum
     field: FieldIdentifierEnum
 
 
-class AggregationSetup(MaxResultsMixin, BaseModel):
+class AggregationSetup(MaxResultsMixin, CustomBaseModel):
     type: AggregationTypeEnum = Field(
         ..., description="Specifies the type of aggregation."
     )
@@ -161,7 +162,7 @@ class AggregationSetup(MaxResultsMixin, BaseModel):
     )
 
 
-class SearchParamsV2(MaxResultsMixin, BaseModel):
+class SearchParamsV2(MaxResultsMixin, CustomBaseModel):
     fields: Optional[List[FieldIdentifierEnum]] = Field(
         default=None,
         description="Fields to include in results. Empty includes all fields.",
@@ -186,10 +187,6 @@ class SearchParamsV2(MaxResultsMixin, BaseModel):
         description="Aggregations to apply for data summarization/analysis.",
     )
 
-    # Forbid extra fields and raise an exception if any are found.
-    class Config:
-        extra = Extra.forbid
-
     @field_validator("sort_order")
     def sort_order_valid(cls, v: Optional[SortOrderEnum]) -> Optional[SortOrderEnum]:
         if v and v not in SortOrderEnum:
diff --git a/audit_logger/models/server_details.py b/audit_logger/models/server_details.py
index 577112d..69be263 100644
--- a/audit_logger/models/server_details.py
+++ b/audit_logger/models/server_details.py
@@ -1,9 +1,11 @@
 from typing import Optional
 
-from pydantic import BaseModel, Field, IPvAnyAddress
+from pydantic import Field, IPvAnyAddress
 
+from audit_logger.models.custom_base import CustomBaseModel
 
-class ServerDetails(BaseModel):
+
+class ServerDetails(CustomBaseModel):
     hostname: Optional[str] = Field(
         default=None,
         description="Hostname of the server where the event occurred.",
diff --git a/audit_logger/utils.py b/audit_logger/utils.py
index 94e71ff..d5de9dd 100644
--- a/audit_logger/utils.py
+++ b/audit_logger/utils.py
@@ -8,14 +8,14 @@
 from elasticsearch import Elasticsearch, SerializationError, helpers
 from faker import Faker
 from fastapi import HTTPException
-from pydantic import HttpUrl, ValidationError
+from pydantic import ValidationError
 
 from audit_logger.custom_logger import get_logger
 from audit_logger.models import (
     ActorDetails,
     AuditLogEntry,
+    BulkAuditLogOptions,
     GenericResponse,
-    RandomAuditLogSettings,
     ResourceDetails,
 )
 from audit_logger.models.env_vars import EnvVars
@@ -98,7 +98,9 @@ def is_valid_ip_v4_address(ip_address: str) -> bool:
         return False
 
 
-def create_bulk_operations(index_name: str, log_entries: List[Dict]) -> List[Dict]:
+def create_bulk_operations(
+    index_name: str, log_entries: List[AuditLogEntry]
+) -> List[Dict]:
     """
     This bulk helper function prepares a list of operations for the Elasticsearch bulk API
     based on the provided log entries.
@@ -129,7 +131,7 @@ def create_bulk_operations(index_name: str, log_entries: List[Dict]) -> List[Dic
 
 
 def generate_audit_log_entries_with_fake_data(
-    settings: RandomAuditLogSettings,
+    settings: BulkAuditLogOptions,
 ) -> List[Dict]:
     """
     Generates a random audit log entry using the Faker library.
@@ -137,12 +139,12 @@ def generate_audit_log_entries_with_fake_data(
     Returns:
     - List[Dict]: A list of audit log entries with fake data.
     """
-    return [generate_fake_log_entry().dict() for _ in range(settings.bulk_count)]
+    return [generate_log_entry().dict() for _ in range(settings.bulk_count)]
 
 
 # ) -> GenericResponse:
 async def process_audit_logs(
-    elastic: Elasticsearch, elastic_index_name: str, log_entries: List[Dict]
+    elastic: Elasticsearch, elastic_index_name: str, log_entries: List[AuditLogEntry]
 ) -> Any:
     """
     Processes a list of audit log entries by sending them to Elasticsearch using the bulk API.
@@ -150,7 +152,7 @@ async def process_audit_logs(
     Args:
     - elastic: An instance of the Elasticsearch client.
     - elastic_index_name (str): The name of the Elasticsearch index.
-    - log_entries (List[Dict]): A list of audit log entries to be processed.
+    - log_entries (List[AuditLogEntry]): A list of audit log entries to be processed.
 
     Returns:
     - GenericResponse
@@ -207,7 +209,7 @@ def validate_date(date_str: str) -> bool:
     return False
 
 
-def generate_fake_log_entry() -> AuditLogEntry:
+def generate_log_entry() -> AuditLogEntry:
     """Create a fake audit log entry using the Faker library."""
     return AuditLogEntry(
         timestamp=fake.date_time_this_year().isoformat(),
@@ -246,10 +248,6 @@ def generate_fake_log_entry() -> AuditLogEntry:
     )
 
 
-def httpurl_to_str(url: HttpUrl) -> str:
-    return str(url)
-
-
 def load_env_vars() -> EnvVars:
     """
     Load all environment variables and validate them against the EnvVars Pydantic model.