diff --git a/lib/cuckoo/core/data/machines.py b/lib/cuckoo/core/data/machines.py index 6f6127f38a3..346adc50e43 100644 --- a/lib/cuckoo/core/data/machines.py +++ b/lib/cuckoo/core/data/machines.py @@ -23,12 +23,15 @@ Select, String, ) + from sqlalchemy.dialects.postgresql import JSONB + from sqlalchemy.ext.mutable import MutableDict from sqlalchemy.orm import ( Mapped, mapped_column, relationship, selectinload, ) + from sqlalchemy.types import JSON except ImportError: # pragma: no cover raise CuckooDependencyError("Unable to import sqlalchemy (install with `poetry install`)") @@ -63,6 +66,13 @@ class Machine(Base): resultserver_port: Mapped[str] = mapped_column(String(255), nullable=False) reserved: Mapped[bool] = mapped_column(Boolean(), nullable=False, default=False) + # Generic key/value bag for per-machine semi-structured data + attributes: Mapped[Optional[dict]] = mapped_column( + MutableDict.as_mutable(JSON().with_variant(JSONB(), "postgresql")), + nullable=True, + default=dict, + ) + def __repr__(self): return f"" @@ -88,6 +98,27 @@ def to_json(self): """ return json.dumps(self.to_dict()) + # ---- attributes helpers ---- + + def get_attribute(self, key: str, default=None): + """Return an attribute value, or default if not set.""" + if not self.attributes: + return default + return self.attributes.get(key, default) + + def set_attribute(self, key: str, value) -> None: + """Set an attribute. Passing None removes the key. Lazy-inits the dict. + + set_attribute is convenience wrapper - MutableDict tracks changes automatically so + direct in-place mutation (machine.attributes["k"] = v) is also fine. + """ + if self.attributes is None: + self.attributes = {} + if value is None: + self.attributes.pop(key, None) + else: + self.attributes[key] = value + def __init__(self, name, label, arch, ip, platform, interface, snapshot, resultserver_ip, resultserver_port, reserved): self.name = name self.label = label diff --git a/lib/cuckoo/core/database.py b/lib/cuckoo/core/database.py index 9b7b9b5e3a7..0582092c71d 100644 --- a/lib/cuckoo/core/database.py +++ b/lib/cuckoo/core/database.py @@ -49,7 +49,7 @@ -SCHEMA_VERSION = "3a1b_tenant_visibility" +SCHEMA_VERSION = "4a6c2b_machine_attributes" log = logging.getLogger(__name__) conf = Config("cuckoo") diff --git a/tests/test_database.py b/tests/test_database.py index 0da02ba471c..0722df9be8d 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -409,6 +409,7 @@ def test_add_machine(self, db: _Database): "resultserver_port": "2043", "arch": "x64", "reserved": False, + "attributes": {}, } assert m2.to_dict() == { @@ -428,8 +429,38 @@ def test_add_machine(self, db: _Database): "tags": ["tag1tag2"], "arch": "x64", "reserved": True, + "attributes": {}, } + def test_machine_attributes(self, db: _Database): + # 1) default=dict — each new row gets its own dict, not shared. + # Flush to apply the default (default fires on flush, not on construction + # for mutable/JSON columns); verify both rows get distinct dict objects. + with db.session.begin(): + m1 = self.add_machine(db, name="attr1", label="attr1") + m2 = self.add_machine(db, name="attr2", label="attr2") + db.session.flush() + assert m1.attributes is not m2.attributes + + # 3) get_attribute / set_attribute helpers. + with db.session.begin(): + m1.set_attribute("last_error", "boom") + m1.set_attribute("vnc_url", "vnc://x") + assert m1.get_attribute("last_error") == "boom" + assert m1.get_attribute("vnc_url") == "vnc://x" + m1.set_attribute("temp", "x") + m1.set_attribute("temp", None) # value=None removes the key + assert m1.get_attribute("temp") is None + + # 2) MutableDict auto-tracks in-place mutation — no flag_modified needed. + with db.session.begin(): + m = db.view_machine("attr1") + m.attributes["direct"] = "via_inplace_mutation" # no set_attribute, no flag_modified + with db.session.begin(): + m = db.view_machine("attr1") + assert m.get_attribute("last_error") == "boom" + assert m.get_attribute("direct") == "via_inplace_mutation" + def test_find_machine_to_service_task_tags_reserved(self, db: _Database): with db.session.begin(): self.add_machine(db, name="name0", label="label0", tags="tag1,x64", reserved=False) diff --git a/utils/db_migration/versions/4. add_machine_attributes.py b/utils/db_migration/versions/4. add_machine_attributes.py new file mode 100644 index 00000000000..2e3ae542e73 --- /dev/null +++ b/utils/db_migration/versions/4. add_machine_attributes.py @@ -0,0 +1,47 @@ +# Copyright (C) 2010-2015 Cuckoo Foundation. +# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org +# See the file 'docs/LICENSE' for copying permission. + +"""Add attributes JSON column to machines + +Revision ID: 4a6c2b_machine_attributes +Revises: 3a1b_tenant_visibility +Create Date: 2026-09-01 + +Adds a generic JSONB column to the machines table to store extra per-machine data - whether + generic (e.g. last shutdown date) or machinery-specific (e.g. VNC connection details) +semi-structured data (e.g. provider metadata, admin notes, runtime state) - without + requiring ALTER TABLE for each new field. + +On Postgres the column is stored as real JSONB; on SQLite/MySQL it falls +back to plain JSON. Matches the Machine model which declares it with +JSON().with_variant(JSONB(), "postgresql") so fresh installs on Postgres +get the optimal type automatically via Base.metadata.create_all(). +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects.postgresql import JSONB + +# revision identifiers, used by Alembic. +revision = "4a6c2b_machine_attributes" +down_revision = "3a1b_tenant_visibility" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # Real JSONB on Postgres (matches the Machine model), plain JSON elsewhere. + # Alembic compiles with_variant against the connection dialect at runtime. + op.add_column( + "machines", + sa.Column( + "attributes", + sa.JSON().with_variant(JSONB(), "postgresql"), + nullable=True, + ), + ) + + +def downgrade() -> None: + op.drop_column("machines", "attributes")