Made changes to db/ and enumerate module __init__ to prep for ZODB transition

pwncat/db/__init__.py

@@ -1,8 +1,5 @@
 #!/usr/bin/env python3
 
-from sqlalchemy.engine import Engine, create_engine
-from sqlalchemy.orm import Session, sessionmaker
-
 import pwncat
 from pwncat.db.base import Base
 from pwncat.db.binary import Binary
@@ -13,53 +10,3 @@ from pwncat.db.suid import SUID
 from pwncat.db.tamper import Tamper
 from pwncat.db.user import User, Group, SecondaryGroupAssociation
 from pwncat.db.fact import Fact
-
-ENGINE: Engine = None
-SESSION_MAKER = None
-SESSION: Session = None
-
-
-def get_engine() -> Engine:
-    """
-    Get a copy of the database engine
-    """
-
-    global ENGINE
-
-    if ENGINE is not None:
-        return ENGINE
-
-    ENGINE = create_engine(pwncat.config["db"], echo=False)
-    Base.metadata.create_all(ENGINE)
-
-    return ENGINE
-
-
-def get_session() -> Session:
-    """
-    Get a new session object
-    """
-
-    global SESSION_MAKER
-    global SESSION
-
-    if SESSION_MAKER is None:
-        SESSION_MAKER = sessionmaker(bind=get_engine())
-    if SESSION is None:
-        SESSION = SESSION_MAKER()
-
-    return SESSION
-
-
-def reset_engine():
-    """
-    Reload the engine and session
-    """
-
-    global ENGINE
-    global SESSION
-    global SESSION_MAKER
-
-    ENGINE = None
-    SESSION = None
-    SESSION_MAKER = None

pwncat/db/base.py

@@ -1,5 +1,3 @@
 #!/usr/bin/env python3
 
-from sqlalchemy.ext.declarative import declarative_base
+# this file is no longer necessary since we are no longer use sqlalchemy?
-
-Base = declarative_base()

pwncat/db/binary.py

@@ -1,18 +1,17 @@
 #!/usr/bin/env python3
-from sqlalchemy import Column, Integer, String, ForeignKey
-from sqlalchemy.orm import relationship
 
-from pwncat.db.base import Base
+import persistent
+from typing import Optional
 
 
-class Binary(Base):
+class Binary(persistent.Persistent):
+    """
+    Stores an understanding of a binary on the target.
+    """
 
-    __tablename__ = "binary"
+    def __init__(self, name, path):
 
-    id = Column(Integer, primary_key=True)
+        # Name of the binary (parameter to which)
-    host_id = Column(Integer, ForeignKey("host.id"))
+        self.name: Optional[str] = name
-    host = relationship("Host", back_populates="binaries")
+        # The path to the binary on the remote host
-    # Name of the binary (parameter to which)
+        self.path: Optional[str] = path
-    name = Column(String)
-    # The path to the binary on the remote host
-    path = Column(String)

pwncat/db/fact.py

@@ -1,36 +1,25 @@
 #!/usr/bin/env python3
-from sqlalchemy import Column, Integer, ForeignKey, PickleType, UniqueConstraint, String
-from sqlalchemy.orm import relationship
 
-from pwncat.db.base import Base
+import persistent
-from pwncat.modules import Result
+from typing import Optional
 
 
-class Fact(Base, Result):
+class Fact(persistent.Persistent):
-    """ Store enumerated facts. The pwncat.enumerate.Fact objects are pickled and
+    """Store enumerated facts. The pwncat.enumerate.Fact objects are pickled and
     stored in the "data" column. The enumerator is arbitrary, but allows for
-    organizations based on the source enumerator. """
+    organizations based on the source enumerator."""
 
-    __tablename__ = "facts"
+    def __init__(self, arg_type, source):
 
-    id = Column(Integer, primary_key=True)
+        # The type of fact (e.g.., "system.user")
-    host_id = Column(Integer, ForeignKey("host.id"))
+        self.type: Optional[str] = arg_type
-    host = relationship("Host", back_populates="facts")
+        # The original procedure that found this fact
-    type = Column(String)
+        self.source: Optional[str] = source
-    source = Column(String)
+
-    data = Column(PickleType)
+        # The original SQLAlchemy-style code held a property, "data",
-    __table_args__ = (
+        # which was a pickle object. We will re-implement that as a subclass
-        UniqueConstraint("type", "data", "host_id", name="_type_data_uc"),
+        # but that may need to include the class properties used previously.
-    )
 
     @property
     def category(self) -> str:
         return f"{self.type}"
-
-    @property
-    def title(self) -> str:
-        return str(self.data)
-
-    @property
-    def description(self) -> str:
-        return getattr(self.data, "description", None)

pwncat/db/history.py

@@ -1,15 +1,13 @@
 #!/usr/bin/env python3
-from sqlalchemy import Column, Integer, String, ForeignKey
-from sqlalchemy.orm import relationship
 
-from pwncat.db.base import Base
+import persistent
+from typing import Optional
 
 
-class History(Base):
+class History(persistent.Persistent):
+    """Store history of ran commands on the target."""
 
-    __tablename__ = "history"
+    def __init__(self, command):
 
-    id = Column(Integer, primary_key=True)
+        # The command ran on the target (e.g., "whoami")
-    host_id = Column(Integer, ForeignKey("host.id"))
+        self.command: Optional[str] = command
-    host = relationship("Host", back_populates="history")
-    command = Column(String)

pwncat/db/host.py

@@ -1,50 +1,4 @@
 #!/usr/bin/env python3
-from sqlalchemy import Column, Integer, String, Enum, Boolean
-from sqlalchemy.orm import relationship
 
-from pwncat import util
+# this file is no longer necessary since this will be encapsulated in the
-from pwncat.db.base import Base
+# target.py file as part of the initial Target object
-
-
-class Host(Base):
-
-    __tablename__ = "host"
-
-    # Database identifier
-    id = Column(Integer, primary_key=True)
-    # A unique hash identifying this host
-    hash = Column(String)
-    # The platform this host is running
-    platform = Column(String)
-    # The IP address we observed on the last connection
-    # to this host
-    ip = Column(String)
-    # The remote architecture (uname -m)
-    arch = Column(String)
-    # The remote init system being used
-    init = Column(Enum(util.Init))
-    # The remote kernel version (uname -r)
-    kernel = Column(String)
-    # The remote distro (probed from /etc/*release), or "unknown"
-    distro = Column(String)
-    # The path to the remote busybox, if installed
-    busybox = Column(String)
-    # Did we install busybox?
-    busybox_uploaded = Column(Boolean)
-    # A list of groups this host has
-    groups = relationship("Group")
-    # A list of users this host has
-    users = relationship("User")
-    # A list of persistence methods applied to this host
-    persistence = relationship("Persistence")
-    # A list of tampers applied to this host
-    tampers = relationship("Tamper")
-    # A list of resolved binaries for the remote host
-    binaries = relationship("Binary")
-    # Command history for local prompt
-    history = relationship("History")
-    # A list of SUID binaries found across all users (may have overlap, and may not be
-    # accessible by the current user).
-    suid = relationship("SUID")
-    # List of enumerated facts about the remote victim
-    facts = relationship("Fact")

pwncat/db/persist.py

@@ -1,21 +1,24 @@
 #!/usr/bin/env python3
-from sqlalchemy import Column, Integer, String, ForeignKey, PickleType
-from sqlalchemy.orm import relationship
 
-from pwncat.db.base import Base
+import persistent
+from typing import Optional
 
 
-class Persistence(Base):
+class Persistence(persistent.Persistent):
+    """
+    Stores an abstract understanding of persistence method installed on a
+    target.
+    """
 
-    __tablename__ = "persistence"
+    def __init__(self, method, user):
 
-    id = Column(Integer, primary_key=True)
+        # The type of persistence
-    host_id = Column(Integer, ForeignKey("host.id"))
+        self.method: Optional[str] = method
-    host = relationship("Host", back_populates="persistence")
+        # The user this persistence was applied as
-    # The type of persistence
+        # (ignored for system persistence)
-    method = Column(String)
+        self.user: Optional[str] = user
-    # The user this persistence was applied as (ignored for system persistence)
+
-    user = Column(String)
+        # The original SQLAlchemy-style code held a property, "args",
-    # The custom arguments passed to the persistence module
+        # which was a pickle object contained the custom arguments passed to
-    # this **will** include the `user` argument.
+        # the persistence module. It **will** include the `user` argument.
-    args = Column(PickleType)
+        # We may re-implement that as a subclass.

pwncat/db/suid.py

@@ -1,20 +1,19 @@
 #!/usr/bin/env python3
-from sqlalchemy import ForeignKey, Integer, Column, String
-from sqlalchemy.orm import relationship
 
-from pwncat.db.base import Base
+import persistent
+from typing import Optional
 
 
-class SUID(Base):
+class SUID(persistent.Persistent):
+    """
+    Stores a record of SUID binaries discovered on the target.
+    """
 
-    __tablename__ = "suid"
+    def __init__(self, path, user):
 
-    id = Column(Integer, primary_key=True)
+        # Path to this SUID binary
-    host_id = Column(Integer, ForeignKey("host.id"))
+        self.path: Optional[str] = path
-    host = relationship("Host", back_populates="suid", foreign_keys=[host_id])
+
-    user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
+        # The original SQLAlchemy-style code held a property, "owner_id",
-    # user = relationship("User", backref="suid", foreign_keys=[user_id])
+        # which maintained the uid corresponding to the user owning this suid
-    # Path to this SUID binary
+        # file. This may or may not be needed?
-    path = Column(String)
-    owner_id = Column(Integer, ForeignKey("users.id"), nullable=False)
-    # owner = relationship("User", foreign_keys=[owner_id], backref="owned_suid")

pwncat/db/tamper.py

@@ -1,16 +1,18 @@
 #!/usr/bin/env python3
-from sqlalchemy import Column, Integer, String, ForeignKey, LargeBinary
-from sqlalchemy.orm import relationship
 
-from pwncat.db.base import Base
+import persistent
+from typing import Optional
 
 
-class Tamper(Base):
+class Tamper(persistent.Persistent):
+    """
+    Stores a record of changes on the target (i.e., things that have been
+    tampered with)
+    """
 
-    __tablename__ = "tamper"
+    def __init__(self, name, data):
 
-    id = Column(Integer, primary_key=True)
+        # The name of this tamper method (what was done on the target)
-    host_id = Column(Integer, ForeignKey("host.id"))
+        self.name: Optional[str] = name
-    host = relationship("Host", back_populates="tampers")
+        # The process outlined in this tamper method
-    name = Column(String)
+        self.data: Optional[bytes] = data
-    data = Column(LargeBinary)

pwncat/db/user.py

@@ -1,67 +1,38 @@
 #!/usr/bin/env python3
-from sqlalchemy import Column, Integer, String, ForeignKey, Table
-from sqlalchemy.orm import relationship
 
-from pwncat.db.base import Base
+import persistent
-
+import persistent.list
-SecondaryGroupAssociation = Table(
+from typing import Optional
-    "secondary_group_association",
-    Base.metadata,
-    Column("group_id", Integer, ForeignKey("groups.id")),
-    Column("user_id", ForeignKey("users.id")),
-)
 
 
-class Group(Base):
+class Group(persistent.Persistent):
+    """
+    Stores a record of changes on the target (i.e., things that have been
+    tampered with)
+    """
 
-    __tablename__ = "groups"
+    def __init__(self, name, members):
 
-    id = Column(Integer, primary_key=True)
+        self.name: Optional[str] = name
-    host_id = Column(Integer, ForeignKey("host.id"), primary_key=True)
+        self.members: persistent.list.PersistentList = persistent.list.PersistentList()
-    host = relationship("Host", back_populates="groups")
-    name = Column(String)
-    members = relationship(
-        "User",
-        back_populates="groups",
-        secondary=SecondaryGroupAssociation,
-        lazy="selectin",
-    )
 
     def __repr__(self):
         return f"""Group(gid={self.id}, name={repr(self.name)}), members={repr(",".join(m.name for m in self.members))})"""
 
 
-class User(Base):
+class User(persistent.Persistent):
+    def __init__(self, name, gid, fullname, homedir, password, hash, shell, groups):
 
-    __tablename__ = "users"
+        self.name: Optional[str] = name
-
+        self.gid: Optional[int] = gid
-    # The users UID
+        self.fullname: Optional[str] = fullname
-    id = Column(Integer, primary_key=True)
+        self.homedir: Optional[str] = homedir
-    host_id = Column(Integer, ForeignKey("host.id"), primary_key=True)
+        self.password: Optional[str] = password
-    host = relationship("Host", back_populates="users", lazy="selectin")
+        self.hash: Optional[str] = hash
-    # The users GID
+        self.shell: Optional[str] = shell
-    gid = Column(Integer, ForeignKey("groups.id"))
+        self.groups: persistent.list.PersistentList = persistent.list.PersistentList(
-    # The actual DB Group object representing that group
+            groups
-    group = relationship("Group")
+        )
-    # The name of the user
-    name = Column(String, primary_key=True)
-    # The user's full name
-    fullname = Column(String)
-    # The user's home directory
-    homedir = Column(String)
-    # The user's password, if known
-    password = Column(String)
-    # The hash of the user's password, if known
-    hash = Column(String)
-    # The user's default shell
-    shell = Column(String)
-    # The user's secondary groups
-    groups = relationship(
-        "Group",
-        back_populates="members",
-        secondary=SecondaryGroupAssociation,
-        lazy="selectin",
-    )
 
     def __repr__(self):
         return f"""User(uid={self.id}, gid={self.gid}, name={repr(self.name)})"""

pwncat/modules/agnostic/enumerate/__init__.py

@@ -3,7 +3,7 @@ from enum import Enum, auto
 import fnmatch
 import time
 
-import sqlalchemy
+# import sqlalchemy
 
 import pwncat
 from pwncat.platform.linux import Linux
@@ -12,7 +12,7 @@ from pwncat.db import get_session
 
 
 class Schedule(Enum):
-    """ Defines how often an enumeration module will run """
+    """Defines how often an enumeration module will run"""
 
     ALWAYS = auto()
     PER_USER = auto()
@@ -20,7 +20,7 @@ class Schedule(Enum):
 
 
 class EnumerateModule(BaseModule):
-    """ Base class for all enumeration modules """
+    """Base class for all enumeration modules"""
 
     # List of categories/enumeration types this module provides
     # This should be set by the sub-classes to know where to find
@@ -49,7 +49,7 @@ class EnumerateModule(BaseModule):
     }
 
     def run(self, session, types, clear):
-        """ Locate all facts this module provides.
+        """Locate all facts this module provides.
 
         Sub-classes should not override this method. Instead, use the
         enumerate method. `run` will cross-reference with database and
@@ -64,34 +64,31 @@ class EnumerateModule(BaseModule):
 
             if clear:
                 # Delete enumerated facts
-                query = db.query(pwncat.db.Fact).filter_by(
+                session.target.facts = persistent.list.PersistentList(
-                    source=self.name, host_id=session.host
+                    (f for f in session.target.facts if f.source != self.name)
                 )
-                query.delete(synchronize_session=False)
+
                 # Delete our marker
-                if self.SCHEDULE != Schedule.ALWAYS:
+                #### We aren't positive how to recreate this in ZODB yet
-                    query = (
+                # if self.SCHEDULE != Schedule.ALWAYS:
-                        db.query(pwncat.db.Fact)
+                #     query = (
-                        .filter_by(host_id=session.host, type="marker")
+                #         db.query(pwncat.db.Fact)
-                        .filter(pwncat.db.Fact.source.startswith(self.name))
+                #         .filter_by(host_id=session.host, type="marker")
-                    )
+                #         .filter(pwncat.db.Fact.source.startswith(self.name))
-                    query.delete(synchronize_session=False)
+                #     )
+                #     query.delete(synchronize_session=False)
                 return
 
             # Yield all the know facts which have already been enumerated
-            existing_facts = (
+            existing_facts = (f for f in session.target.facts if f.source == self.name)
-                db.query(pwncat.db.Fact)
-                .filter_by(source=self.name, host_id=session.host)
-                .filter(pwncat.db.Fact.type != "marker")
-            )
 
             if types:
-                for fact in existing_facts.all():
+                for fact in existing_facts:
                     for typ in types:
                         if fnmatch.fnmatch(fact.type, typ):
                             yield fact
             else:
-                yield from existing_facts.all()
+                yield from existing_facts
 
             if self.SCHEDULE != Schedule.ALWAYS:
                 exists = (
@@ -110,10 +107,11 @@ class EnumerateModule(BaseModule):
                     continue
 
                 typ, data = item
+                # session.target.facts.append(fact)
 
-                row = pwncat.db.Fact(
+                # row = pwncat.db.Fact(
-                    host_id=session.host, type=typ, data=data, source=self.name
+                #     host_id=session.host, type=typ, data=data, source=self.name
-                )
+                # )
                 try:
                     db.add(row)
                     db.commit()
@@ -135,13 +133,19 @@ class EnumerateModule(BaseModule):
             # Add the marker if needed
             if self.SCHEDULE != Schedule.ALWAYS:
                 row = pwncat.db.Fact(
-                    host_id=session.host, type="marker", source=marker_name, data=None,
+                    host_id=session.host,
+                    type="marker",
+                    source=marker_name,
+                    data=None,
                 )
                 db.add(row)
+            # session.db.transaction_manager.commit()
 
     def enumerate(self, session):
-        """ Defined by sub-classes to do the actual enumeration of
+        """
-        facts. """
+        Defined by sub-classes to do the actual enumeration of
+        facts.
+        """
 
 
 # This makes `run enumerate` initiate a quick scan

pwncat/modules/agnostic/enumerate/gather.py

@@ -22,7 +22,7 @@ def strip_markup(styled_text: str) -> str:
 
 
 def list_wrapper(iterable):
-    """ Wraps a list in a generator """
+    """Wraps a list in a generator"""
     yield from iterable
 
 
@@ -72,7 +72,7 @@ class Module(pwncat.modules.BaseModule):
     PLATFORM = None
 
     def run(self, session, output, modules, types, clear):
-        """ Perform a enumeration of the given moduels and save the output """
+        """Perform a enumeration of the given moduels and save the output"""
 
         module_names = modules
 

pwncat/modules/agnostic/enumerate/quick.py

@@ -5,7 +5,7 @@ from pwncat.modules import BaseModule, Status, Argument
 
 
 class Module(BaseModule):
-    """ Perform a quick enumeration of common useful data """
+    """Perform a quick enumeration of common useful data"""
 
     ARGUMENTS = {
         "output": Argument(