From d85dbdd0b48de78d523c50a2d00dce7568ab0caf Mon Sep 17 00:00:00 2001 From: John Hammond Date: Fri, 30 Apr 2021 21:34:30 -0400 Subject: [PATCH] Made changes to db/ and enumerate module __init__ to prep for ZODB transition --- pwncat/db/__init__.py | 53 ------------- pwncat/db/base.py | 4 +- pwncat/db/binary.py | 23 +++--- pwncat/db/fact.py | 39 ++++------ pwncat/db/history.py | 16 ++-- pwncat/db/host.py | 50 +------------ pwncat/db/persist.py | 33 ++++---- pwncat/db/suid.py | 27 ++++--- pwncat/db/tamper.py | 22 +++--- pwncat/db/user.py | 75 ++++++------------- pwncat/modules/agnostic/enumerate/__init__.py | 58 +++++++------- pwncat/modules/agnostic/enumerate/gather.py | 4 +- pwncat/modules/agnostic/enumerate/quick.py | 2 +- 13 files changed, 135 insertions(+), 271 deletions(-) diff --git a/pwncat/db/__init__.py b/pwncat/db/__init__.py index b59ec35..6f9616c 100644 --- a/pwncat/db/__init__.py +++ b/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 diff --git a/pwncat/db/base.py b/pwncat/db/base.py index 992602f..80c1522 100644 --- a/pwncat/db/base.py +++ b/pwncat/db/base.py @@ -1,5 +1,3 @@ #!/usr/bin/env python3 -from sqlalchemy.ext.declarative import declarative_base - -Base = declarative_base() +# this file is no longer necessary since we are no longer use sqlalchemy? diff --git a/pwncat/db/binary.py b/pwncat/db/binary.py index 47896ab..4caab86 100644 --- a/pwncat/db/binary.py +++ b/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) - host_id = Column(Integer, ForeignKey("host.id")) - host = relationship("Host", back_populates="binaries") - # Name of the binary (parameter to which) - name = Column(String) - # The path to the binary on the remote host - path = Column(String) + # Name of the binary (parameter to which) + self.name: Optional[str] = name + # The path to the binary on the remote host + self.path: Optional[str] = path diff --git a/pwncat/db/fact.py b/pwncat/db/fact.py index 7d92c9a..b71ba79 100644 --- a/pwncat/db/fact.py +++ b/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 -from pwncat.modules import Result +import persistent +from typing import Optional -class Fact(Base, Result): - """ Store enumerated facts. The pwncat.enumerate.Fact objects are pickled and +class Fact(persistent.Persistent): + """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) - host_id = Column(Integer, ForeignKey("host.id")) - host = relationship("Host", back_populates="facts") - type = Column(String) - source = Column(String) - data = Column(PickleType) - __table_args__ = ( - UniqueConstraint("type", "data", "host_id", name="_type_data_uc"), - ) + # The type of fact (e.g.., "system.user") + self.type: Optional[str] = arg_type + # The original procedure that found this fact + self.source: Optional[str] = source + + # The original SQLAlchemy-style code held a property, "data", + # which was a pickle object. We will re-implement that as a subclass + # 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) diff --git a/pwncat/db/history.py b/pwncat/db/history.py index 303fb3d..be3601c 100644 --- a/pwncat/db/history.py +++ b/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) - host_id = Column(Integer, ForeignKey("host.id")) - host = relationship("Host", back_populates="history") - command = Column(String) + # The command ran on the target (e.g., "whoami") + self.command: Optional[str] = command diff --git a/pwncat/db/host.py b/pwncat/db/host.py index 8599364..a99f6a5 100644 --- a/pwncat/db/host.py +++ b/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 -from pwncat.db.base import Base - - -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") +# this file is no longer necessary since this will be encapsulated in the +# target.py file as part of the initial Target object diff --git a/pwncat/db/persist.py b/pwncat/db/persist.py index 4d1b245..9085161 100644 --- a/pwncat/db/persist.py +++ b/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) - host_id = Column(Integer, ForeignKey("host.id")) - host = relationship("Host", back_populates="persistence") - # The type of persistence - method = Column(String) - # The user this persistence was applied as (ignored for system persistence) - user = Column(String) - # The custom arguments passed to the persistence module - # this **will** include the `user` argument. - args = Column(PickleType) + # The type of persistence + self.method: Optional[str] = method + # The user this persistence was applied as + # (ignored for system persistence) + self.user: Optional[str] = user + + # The original SQLAlchemy-style code held a property, "args", + # which was a pickle object contained the custom arguments passed to + # the persistence module. It **will** include the `user` argument. + # We may re-implement that as a subclass. diff --git a/pwncat/db/suid.py b/pwncat/db/suid.py index 9c20485..c1cea3d 100644 --- a/pwncat/db/suid.py +++ b/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) - host_id = Column(Integer, ForeignKey("host.id")) - host = relationship("Host", back_populates="suid", foreign_keys=[host_id]) - user_id = Column(Integer, ForeignKey("users.id"), nullable=False) - # user = relationship("User", backref="suid", foreign_keys=[user_id]) - # Path to this SUID binary - path = Column(String) - owner_id = Column(Integer, ForeignKey("users.id"), nullable=False) - # owner = relationship("User", foreign_keys=[owner_id], backref="owned_suid") + # Path to this SUID binary + self.path: Optional[str] = path + + # The original SQLAlchemy-style code held a property, "owner_id", + # which maintained the uid corresponding to the user owning this suid + # file. This may or may not be needed? diff --git a/pwncat/db/tamper.py b/pwncat/db/tamper.py index 3c7ef71..d7100d3 100644 --- a/pwncat/db/tamper.py +++ b/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) - host_id = Column(Integer, ForeignKey("host.id")) - host = relationship("Host", back_populates="tampers") - name = Column(String) - data = Column(LargeBinary) + # The name of this tamper method (what was done on the target) + self.name: Optional[str] = name + # The process outlined in this tamper method + self.data: Optional[bytes] = data diff --git a/pwncat/db/user.py b/pwncat/db/user.py index 6f3c765..63ea449 100644 --- a/pwncat/db/user.py +++ b/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 - -SecondaryGroupAssociation = Table( - "secondary_group_association", - Base.metadata, - Column("group_id", Integer, ForeignKey("groups.id")), - Column("user_id", ForeignKey("users.id")), -) +import persistent +import persistent.list +from typing import Optional -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) - host_id = Column(Integer, ForeignKey("host.id"), primary_key=True) - host = relationship("Host", back_populates="groups") - name = Column(String) - members = relationship( - "User", - back_populates="groups", - secondary=SecondaryGroupAssociation, - lazy="selectin", - ) + self.name: Optional[str] = name + self.members: persistent.list.PersistentList = persistent.list.PersistentList() 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" - - # The users UID - id = Column(Integer, primary_key=True) - host_id = Column(Integer, ForeignKey("host.id"), primary_key=True) - host = relationship("Host", back_populates="users", lazy="selectin") - # The users GID - gid = Column(Integer, ForeignKey("groups.id")) - # The actual DB Group object representing that group - 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", - ) + self.name: Optional[str] = name + self.gid: Optional[int] = gid + self.fullname: Optional[str] = fullname + self.homedir: Optional[str] = homedir + self.password: Optional[str] = password + self.hash: Optional[str] = hash + self.shell: Optional[str] = shell + self.groups: persistent.list.PersistentList = persistent.list.PersistentList( + groups + ) def __repr__(self): return f"""User(uid={self.id}, gid={self.gid}, name={repr(self.name)})""" diff --git a/pwncat/modules/agnostic/enumerate/__init__.py b/pwncat/modules/agnostic/enumerate/__init__.py index 75820ce..bd28e25 100644 --- a/pwncat/modules/agnostic/enumerate/__init__.py +++ b/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( - source=self.name, host_id=session.host + session.target.facts = persistent.list.PersistentList( + (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: - query = ( - db.query(pwncat.db.Fact) - .filter_by(host_id=session.host, type="marker") - .filter(pwncat.db.Fact.source.startswith(self.name)) - ) - query.delete(synchronize_session=False) + #### We aren't positive how to recreate this in ZODB yet + # if self.SCHEDULE != Schedule.ALWAYS: + # query = ( + # db.query(pwncat.db.Fact) + # .filter_by(host_id=session.host, type="marker") + # .filter(pwncat.db.Fact.source.startswith(self.name)) + # ) + # query.delete(synchronize_session=False) return # Yield all the know facts which have already been enumerated - existing_facts = ( - db.query(pwncat.db.Fact) - .filter_by(source=self.name, host_id=session.host) - .filter(pwncat.db.Fact.type != "marker") - ) + existing_facts = (f for f in session.target.facts if f.source == self.name) 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( - host_id=session.host, type=typ, data=data, source=self.name - ) + # row = pwncat.db.Fact( + # 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 diff --git a/pwncat/modules/agnostic/enumerate/gather.py b/pwncat/modules/agnostic/enumerate/gather.py index c287034..ebffa63 100644 --- a/pwncat/modules/agnostic/enumerate/gather.py +++ b/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 diff --git a/pwncat/modules/agnostic/enumerate/quick.py b/pwncat/modules/agnostic/enumerate/quick.py index 4c0aec5..0e1c42a 100644 --- a/pwncat/modules/agnostic/enumerate/quick.py +++ b/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(