mmisp-lib 0.1.4__tar.gz → 0.1.6__tar.gz

mmisp_lib-0.1.6/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Konstantin Zangerle
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/PKG-INFO

@@ -1,8 +1,9 @@
 Metadata-Version: 2.1
 Name: mmisp-lib
-Version: 0.1.4
+Version: 0.1.6
 Requires-Python: >=3.11.0
 Description-Content-Type: text/markdown
+License-File: LICENSE
 Requires-Dist: fastapi==0.104.1
 Requires-Dist: SQLAlchemy[asyncio]==1.4.46
 Requires-Dist: pydantic==1.10.13

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mmisp-lib"
-version = "0.1.4"
+version = "0.1.6"
 description = ""
 authors = []
 readme = "README.md"

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/api_schemas/attributes/add_attribute_body.py

@@ -2,11 +2,13 @@ from typing import Any, Optional
 
 from pydantic import BaseModel, root_validator
 
+from mmisp.lib.attributes import literal_valid_attribute_types
+
 
 class AddAttributeBody(BaseModel):
-    type: str
-    value: Optional[str]
-    value1: Optional[str]
+    type: literal_valid_attribute_types
+    value: str | None = None
+    value1: str | None = None
     value2: str | None = None
     event_id: str | None = None
     object_id: str | None = None
@@ -29,6 +31,3 @@ class AddAttributeBody(BaseModel):
         if all(item is None for item in required_values):
             raise ValueError("value or value1 has to be set")
         return data
-
-    class Config:
-        orm_mode = True

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/api_schemas/attributes/add_attribute_response.py

@@ -18,12 +18,12 @@ class AddAttributeAttributes(BaseModel):
     timestamp: str
     distribution: str
     sharing_group_id: str
-    comment: str
+    comment: str | None = None
     deleted: bool
     disable_correlation: bool
-    first_seen: Optional[str] = Field(..., nullable=True)
-    last_seen: Optional[str] = Field(..., nullable=True)
-    attribute_tag: list[str] = Field([], alias="AttributeTag")
+    first_seen: str | None = None
+    last_seen: str | None = None
+    attribute_tag: list[str] | None = Field(default_factory=list, alias="AttributeTag")
 
 
 class AddAttributeResponse(BaseModel):

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/api_schemas/objects/create_object_body.py

@@ -14,6 +14,3 @@ class ObjectCreateBody(BaseModel):
     first_seen: str | None = None
     last_seen: str | None = None
     Attribute: list[AddAttributeBody] | None = None
-
-    class Config:
-        orm_mode = True

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/api_schemas/objects/get_object_response.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, Optional
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, Field, validator
 
 from mmisp.api_schemas.attributes.get_all_attributes_response import GetAllAttributesResponse
 from mmisp.api_schemas.events.get_event_response import ObjectEventResponse
@@ -22,7 +22,7 @@ class ObjectWithAttributesResponse(BaseModel):
     deleted: bool | None = None
     first_seen: str | None = None
     last_seen: str | None = None
-    Attribute: list[GetAllAttributesResponse] | None = None
+    attributes: list[GetAllAttributesResponse] | None = Field(alias="Attribute", default=None)
     Event: ObjectEventResponse | None = None
 
     @validator("sharing_group_id", always=True)
@@ -35,16 +35,13 @@ class ObjectWithAttributesResponse(BaseModel):
             return value
         return None
 
+    class Config:
+        allow_population_by_field_name = True
+
 
 class ObjectResponse(BaseModel):
     Object: ObjectWithAttributesResponse
 
-    class Config:
-        orm_mode = True
-
 
 class ObjectSearchResponse(BaseModel):
     response: list[ObjectResponse]
-
-    class Config:
-        orm_mode = True

mmisp_lib-0.1.6/src/mmisp/db/all_models.py

@@ -0,0 +1,41 @@
+from mmisp.db.models import (  # noqa: F401
+    attribute,
+    auth_key,
+    event,
+    feed,
+    galaxy,
+    identity_provider,
+    noticelist,
+    object,
+    organisation,
+    role,
+    server,
+    sharing_group,
+    sighting,
+    tag,
+    taxonomy,
+    user,
+    user_setting,
+    warninglist,
+)
+
+__all__ = [
+    attribute,
+    auth_key,
+    event,
+    feed,
+    galaxy,
+    identity_provider,
+    noticelist,
+    object,
+    organisation,
+    role,
+    server,
+    sharing_group,
+    sighting,
+    tag,
+    taxonomy,
+    user,
+    user_setting,
+    warninglist,
+]

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/db/config.py

@@ -7,9 +7,10 @@ from dotenv import load_dotenv
 @dataclass
 class DatabaseConfig:
     DATABASE_URL: str
+    DEBUG: bool
 
 
 load_dotenv(getenv("ENV_FILE", ".env"))
 
 
-config: DatabaseConfig = DatabaseConfig(DATABASE_URL=getenv("DATABASE_URL", ""))
+config: DatabaseConfig = DatabaseConfig(DATABASE_URL=getenv("DATABASE_URL", ""), DEBUG=bool(getenv("DEBUG", False)))

mmisp_lib-0.1.6/src/mmisp/db/database.py

@@ -0,0 +1,89 @@
+import contextlib
+from typing import AsyncIterator, Self, TypeAlias
+
+from sqlalchemy.engine.url import make_url
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, AsyncSession, create_async_engine
+from sqlalchemy.orm import declarative_base, sessionmaker
+
+from mmisp.db.config import config
+
+Session: TypeAlias = AsyncSession
+
+# , poolclass=NullPool)
+# async_session = sessionmaker(autoflush=False, expire_on_commit=False, class_=AsyncSession, bind=engine)
+
+Base = declarative_base()
+
+
+class DatabaseSessionManager:
+    def __init__(self: Self) -> None:
+        self._engine: AsyncEngine | None = None
+        self._sessionmaker: sessionmaker | None = None
+
+        self._url = make_url(config.DATABASE_URL)
+
+    def init(self: Self) -> None:
+        if config.DEBUG:
+            self._engine = create_async_engine(self._url, echo=True)
+        else:
+            self._engine = create_async_engine(self._url)
+        self._sessionmaker = sessionmaker(
+            autocommit=False, expire_on_commit=False, bind=self._engine, class_=AsyncSession
+        )
+
+    async def close(self: Self) -> None:
+        if self._engine is None:
+            raise Exception("DatabaseSessionManager is not initialized")
+        await self._engine.dispose()
+        self._engine = None
+        self._sessionmaker = None
+
+    @contextlib.asynccontextmanager
+    async def connect(self: Self) -> AsyncIterator[AsyncConnection]:
+        if self._engine is None:
+            raise Exception("DatabaseSessionManager is not initialized")
+
+        async with self._engine.begin() as connection:
+            try:
+                yield connection
+            except Exception:
+                await connection.rollback()
+                raise
+
+    @contextlib.asynccontextmanager
+    async def session(self: Self) -> AsyncIterator[AsyncSession]:
+        if self._sessionmaker is None:
+            raise Exception("DatabaseSessionManager is not initialized")
+
+        session = self._sessionmaker()
+        try:
+            yield session
+        except Exception:
+            await session.rollback()
+            raise
+        finally:
+            await session.close()
+
+    # Used for testing
+    async def create_all(self: Self, engine: AsyncEngine | None = None) -> None:
+        if engine is None:
+            engine = self._engine
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
+
+    async def drop_all(self: Self, engine: AsyncEngine | None = None) -> None:
+        if engine is None:
+            engine = self._engine
+        await engine.run_sync(Base.metadata.drop_all)
+
+
+async def get_db() -> Session:
+    async with sessionmanager.session() as session:
+        yield session
+
+
+async def create_all_models() -> None:
+    await sessionmanager.create_all()
+
+
+sessionmanager = DatabaseSessionManager()

mmisp_lib-0.1.6/src/mmisp/db/mixins.py

@@ -0,0 +1,17 @@
+from typing import Self
+
+from sqlalchemy import inspect
+from sqlalchemy.ext.hybrid import hybrid_property
+
+
+class DictMixin:
+    def asdict(self: Self) -> dict:
+        d = {}
+        for key in self.__mapper__.c.keys():
+            if not key.startswith("_"):
+                d[key] = getattr(self, key)
+
+        for key, prop in inspect(self.__class__).all_orm_descriptors.items():
+            if isinstance(prop, hybrid_property):
+                d[key] = getattr(self, key)
+        return d

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/db/models/attribute.py

@@ -1,10 +1,11 @@
 from typing import Self, Type
 
-from sqlalchemy import BigInteger, Boolean, Column, ForeignKey, Integer, String, Text, inspect
+from sqlalchemy import BigInteger, Boolean, Column, ForeignKey, Integer, String, Text
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import relationship
 from sqlalchemy.orm.decl_api import DeclarativeMeta
 
+from mmisp.db.mixins import DictMixin
 from mmisp.lib.attributes import categories, default_category, mapper_safe_clsname_val, to_ids
 from mmisp.util.uuid import uuid
 
@@ -13,19 +14,6 @@ from .event import Event
 from .tag import Tag
 
 
-class DictMixin:
-    def asdict(self: Self) -> dict:
-        d = {}
-        for key in self.__mapper__.c.keys():
-            if not key.startswith("_"):
-                d[key] = getattr(self, key)
-
-        for key, prop in inspect(self.__class__).all_orm_descriptors.items():
-            if isinstance(prop, hybrid_property):
-                d[key] = getattr(self, key)
-        return d
-
-
 class Attribute(Base, DictMixin):
     __tablename__ = "attributes"
 
@@ -50,6 +38,17 @@ class Attribute(Base, DictMixin):
 
     event = relationship("Event", back_populates="attributes", lazy="joined")
 
+    __mapper_args__ = {"polymorphic_on": "type"}
+
+    def __init__(self: Self, *arg, **kwargs) -> None:
+        if kwargs["value1"] is None:
+            split_val = kwargs["value"].split("|", 1)
+            kwargs["value1"] = split_val[0]
+            if len(split_val) == 2:
+                kwargs["value2"] = split_val[1]
+
+        super().__init__(*arg, **kwargs)
+
     @property
     def event_uuid(self: "Attribute") -> str:
         return self.event.uuid
@@ -83,7 +82,7 @@ class AttributeMeta(DeclarativeMeta):
         key = clsname[len("Attribute") :]
         dct["default_category"] = default_category[mapper_safe_clsname_val[key]]
         dct["categories"] = categories[mapper_safe_clsname_val[key]]
-        dct["to_ids"] = to_ids[mapper_safe_clsname_val[key]]
+        dct["default_to_ids"] = to_ids[mapper_safe_clsname_val[key]]
         dct["__mapper_args__"] = {"polymorphic_identity": mapper_safe_clsname_val[key]}
         return super().__new__(cls, clsname, bases, dct)
 

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/db/models/object.py

@@ -1,10 +1,11 @@
 from sqlalchemy import Boolean, Column, ForeignKey, Integer, String
 
 from mmisp.db.database import Base
+from mmisp.db.mixins import DictMixin
 from mmisp.util.uuid import uuid
 
 
-class Object(Base):
+class Object(Base, DictMixin):
     __tablename__ = "objects"
 
     id = Column(Integer, primary_key=True, nullable=False)
@@ -24,7 +25,7 @@ class Object(Base):
     last_seen = Column(Integer, index=True, default=None)
 
 
-class ObjectTemplate(Base):
+class ObjectTemplate(Base, DictMixin):
     __tablename__ = "object_templates"
 
     id = Column(Integer, primary_key=True, nullable=False)

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/db/models/organisation.py

@@ -2,12 +2,13 @@ from datetime import datetime
 
 from sqlalchemy import Boolean, Column, DateTime, Integer, String, Text
 
+from mmisp.db.mixins import DictMixin
 from mmisp.util.uuid import uuid
 
 from ..database import Base
 
 
-class Organisation(Base):
+class Organisation(Base, DictMixin):
     __tablename__ = "organisations"
 
     id = Column(Integer, primary_key=True, nullable=False)

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/db/models/sharing_group.py

@@ -2,12 +2,13 @@ from datetime import datetime
 
 from sqlalchemy import Boolean, Column, DateTime, Integer, String, Text
 
+from mmisp.db.mixins import DictMixin
 from mmisp.util.uuid import uuid
 
 from ..database import Base
 
 
-class SharingGroup(Base):
+class SharingGroup(Base, DictMixin):
     __tablename__ = "sharing_groups"
 
     id = Column(Integer, primary_key=True, nullable=False)
@@ -25,7 +26,7 @@ class SharingGroup(Base):
     roaming = Column(Boolean, default=False, nullable=False)
 
 
-class SharingGroupOrg(Base):
+class SharingGroupOrg(Base, DictMixin):
     __tablename__ = "sharing_group_orgs"
 
     id = Column(Integer, primary_key=True, nullable=False)
@@ -34,7 +35,7 @@ class SharingGroupOrg(Base):
     extend = Column(Boolean, default=False, nullable=False)
 
 
-class SharingGroupServer(Base):
+class SharingGroupServer(Base, DictMixin):
     __tablename__ = "sharing_group_servers"
 
     id = Column(Integer, primary_key=True, nullable=False)

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp/lib/attributes.py

@@ -1,8 +1,9 @@
 from collections import defaultdict
-from enum import Enum
+from enum import StrEnum
+from typing import Literal, Union
 
 
-class AttributeCategories(Enum):
+class AttributeCategories(StrEnum):
     PAYLOAD_DELIVERY = "Payload delivery"
     ARTIFACTS_DROPPED = "Artifacts dropped"
     PAYLOAD_INSTALLATION = "Payload installation"
@@ -213,7 +214,7 @@ mapper_val_safe_clsname = {
 }
 
 mapper_safe_clsname_val = dict((v, k) for k, v in mapper_val_safe_clsname.items())
-
+literal_valid_attribute_types = Union[*[Literal[k] for k in mapper_val_safe_clsname.keys()]]
 
 default_category = {
     "md5": AttributeCategories.PAYLOAD_DELIVERY,

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp_lib.egg-info/PKG-INFO

@@ -1,8 +1,9 @@
 Metadata-Version: 2.1
 Name: mmisp-lib
-Version: 0.1.4
+Version: 0.1.6
 Requires-Python: >=3.11.0
 Description-Content-Type: text/markdown
+License-File: LICENSE
 Requires-Dist: fastapi==0.104.1
 Requires-Dist: SQLAlchemy[asyncio]==1.4.46
 Requires-Dist: pydantic==1.10.13

{mmisp_lib-0.1.4 → mmisp_lib-0.1.6}/src/mmisp_lib.egg-info/SOURCES.txt

@@ -1,3 +1,4 @@
+LICENSE
 README.md
 pyproject.toml
 src/mmisp/api_schemas/__init__.py
@@ -132,8 +133,10 @@ src/mmisp/api_schemas/warninglists/toggle_enable_warninglists_body.py
 src/mmisp/api_schemas/warninglists/toggle_enable_warninglists_response.py
 src/mmisp/api_schemas/warninglists/warninglist_response.py
 src/mmisp/db/__init__.py
+src/mmisp/db/all_models.py
 src/mmisp/db/config.py
 src/mmisp/db/database.py
+src/mmisp/db/mixins.py
 src/mmisp/db/print_changes.py
 src/mmisp/db/models/__init__.py
 src/mmisp/db/models/attribute.py

mmisp_lib-0.1.4/src/mmisp/db/database.py

@@ -1,42 +0,0 @@
-from functools import wraps
-from typing import Any, Callable, TypeAlias
-
-from sqlalchemy.engine.url import make_url
-from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
-from sqlalchemy.orm import declarative_base, sessionmaker
-from sqlalchemy.pool import NullPool
-
-from mmisp.db.config import config
-
-Session: TypeAlias = AsyncSession
-
-url = make_url(config.DATABASE_URL)
-engine = create_async_engine(url, poolclass=NullPool)
-async_session = sessionmaker(autoflush=False, expire_on_commit=False, class_=AsyncSession, bind=engine)
-
-Base = declarative_base()
-
-
-def get_db() -> Session:
-    return async_session()
-
-
-def with_session_management(fn: Callable) -> Callable:
-    @wraps(fn)
-    async def wrapper(*args, **kwargs) -> Any:
-        db: Session = kwargs.pop("db")
-        output: Any = None
-
-        try:
-            output = await fn(*args, **kwargs, db=db)
-        finally:
-            await db.close()
-
-        return output
-
-    return wrapper
-
-
-async def create_all_models() -> None:
-    async with engine.begin() as db:
-        await db.run_sync(Base.metadata.create_all)