From 3e501d2957ff52b1821b6de838c40e25bf8291c4 Mon Sep 17 00:00:00 2001 From: Caleb Stewart Date: Fri, 18 Jun 2021 22:04:12 -0400 Subject: [PATCH] Added enumeration scopes --- CHANGELOG.md | 1 + pwncat/channel/socket.py | 2 +- pwncat/manager.py | 17 ++- pwncat/modules/agnostic/enumerate/gather.py | 49 +------ pwncat/modules/enumerate.py | 127 +++++++++++++----- .../windows/enumerate/system/processes.py | 5 +- pwncat/platform/windows.py | 1 - 7 files changed, 118 insertions(+), 84 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 559c0fe..7f50910 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ and simply didn't have the time to go back and retroactively create one. - Added `ncat`-style ssl arguments to entrypoint and `connect` command - Added query-string arguments to connection strings for both the entrypoint and the `connect` command. +- Added Enumeration States to allow session-bound enumerations ## [0.4.3] - 2021-06-18 Patch fix release. Major fixes are the correction of file IO for LinuxWriters and diff --git a/pwncat/channel/socket.py b/pwncat/channel/socket.py index ec067d9..e87aee7 100644 --- a/pwncat/channel/socket.py +++ b/pwncat/channel/socket.py @@ -22,7 +22,7 @@ import socket import functools from typing import Optional -from pwncat.channel import Channel, ChannelClosed +from pwncat.channel import Channel, ChannelError, ChannelClosed def connect_required(method): diff --git a/pwncat/manager.py b/pwncat/manager.py index 4354f9a..1e46f55 100644 --- a/pwncat/manager.py +++ b/pwncat/manager.py @@ -41,6 +41,7 @@ from pwncat.target import Target from pwncat.channel import Channel, ChannelError, ChannelClosed from pwncat.commands import CommandParser from pwncat.platform import Platform, PlatformError +from pwncat.modules.enumerate import Scope class InteractiveExit(Exception): @@ -68,6 +69,8 @@ class Session: self.module_depth = 0 self.showing_progress = True self.layers = [] + self.enumerate_state = {} + self.facts = [] self._progress = None @@ -170,15 +173,23 @@ class Session: if members is None or any(m in group.members for m in members): yield group - def register_fact(self, fact: "pwncat.db.Fact"): + def register_fact( + self, + fact: "pwncat.db.Fact", + scope: Scope = Scope.HOST, + commit: bool = False, + ): """Register a fact with this session's target. This is useful when a fact is generated during execution of a command or module, but is not associated with a specific enumeration module. It can still be queried with the base `enumerate` module by it's type.""" - if fact not in self.target.facts: + if scope is Scope.HOST and fact not in self.target.facts: self.target.facts.append(fact) - self.db.transaction_manager.commit() + if commit: + self.db.transaction_manager.commit() + elif scope is Scope.SESSION and fact not in self.facts: + self.facts.append(fact) def run(self, module: str, **kwargs): """Run a module on this session""" diff --git a/pwncat/modules/agnostic/enumerate/gather.py b/pwncat/modules/agnostic/enumerate/gather.py index 1739cdf..2cf2cec 100644 --- a/pwncat/modules/agnostic/enumerate/gather.py +++ b/pwncat/modules/agnostic/enumerate/gather.py @@ -3,12 +3,15 @@ import fnmatch from io import IOBase import pwncat.modules -from pwncat import util -from pwncat.util import strip_markup from pwncat.modules import Status, ModuleFailed from pwncat.modules.enumerate import EnumerateModule +def iterate_two_lists(a, b): + yield from a + yield from b + + def list_wrapper(iterable): """Wraps a list in a generator""" yield from iterable @@ -98,7 +101,7 @@ class Module(pwncat.modules.BaseModule): facts = {} if cache: - for fact in session.target.facts: + for fact in iterate_two_lists(session.target.facts, session.facts): if not types or any( any(fnmatch.fnmatch(t2, t1) for t2 in fact.types) for t1 in types ): @@ -154,43 +157,3 @@ class Module(pwncat.modules.BaseModule): yield Status(item.title(session)) except ModuleFailed as exc: session.log(f"[red]{module.name}[/red]: {str(exc)}") - - # We didn't ask for a report output file, so don't write one. - # Because output is none, the results were already returned - # in the above loop. - if output is None: - return - - yield pwncat.modules.Status("writing report") - - with output as filp: - - with session.db as db: - host = db.query(pwncat.db.Host).filter_by(id=session.host).first() - - filp.write(f"# {host.ip} - Enumeration Report\n\n") - filp.write("Enumerated Types:\n") - for typ in facts: - filp.write(f"- {typ}\n") - filp.write("\n") - - for typ in facts: - - filp.write(f"## {typ.upper()} Facts\n\n") - - sections = [] - for fact in facts[typ]: - if getattr(fact.data, "description", None) is not None: - sections.append(fact) - continue - filp.write( - f"- {util.escape_markdown(strip_markup(str(fact.data)))}\n" - ) - - filp.write("\n") - - for section in sections: - filp.write( - f"### {util.escape_markdown(strip_markup(str(section.data)))}\n\n" - ) - filp.write(f"```\n{section.data.description}\n```\n\n") diff --git a/pwncat/modules/enumerate.py b/pwncat/modules/enumerate.py index bc72f32..1a22706 100644 --- a/pwncat/modules/enumerate.py +++ b/pwncat/modules/enumerate.py @@ -18,6 +18,15 @@ requested. Unlike base modules, enumeration modules do not accept any custom arguments. However, they do still require a list of compatible platforms. +Aside from schedules, you can also specify an enumeration scope. The default +scope is ``Scope.HOST``. This scope saves facts to the database which is +shared between sessions and between instances of pwncat. ``Scope.SESSION`` +defines facts which live only as long as the specific session is alive and +are not shared with other sessions with the same target. The ``Scope.NONE`` +scope specifies that facts are never saved. This is normally used with +``Schedule.ALWAYS`` to have enumeration modules run every time without saving +the facts. + When defining an enumeration module, you must define the :func:`EnumerateModule.enumerate` method. This method is a generator which can yield either facts or status updates, just like the @@ -44,6 +53,7 @@ Example Enumerate Module PLATFORM = [Windows] SCHEDULE = Schedule.PER_USER PROVIDES = ["custom.fact.type"] + SCOPE = Scope.HOST def enumerate(self, session: "pwncat.manager.Session"): yield CustomFactObject(self.name) @@ -53,8 +63,6 @@ import typing import fnmatch from enum import Enum, auto -import persistent - import pwncat from pwncat.db import Fact from pwncat.modules import List, Status, Argument, BaseModule @@ -66,14 +74,25 @@ class Schedule(Enum): ALWAYS = auto() """ Execute the enumeration every time the module is executed """ - NOSAVE = auto() - """ Similar to always, except that no enumerated information will be saved to the database """ PER_USER = auto() """ Execute the enumeration once per user on the target """ ONCE = auto() """ Execute the enumeration once and only once """ +class Scope(Enum): + """Defines whether the fact is scoped to the target host or + to the active session. Session-scoped facts are lost when a + session ends.""" + + HOST = auto() + """ Host scope; facts are saved in the database """ + SESSION = auto() + """ Session scope; facts are lost when the session ends """ + NONE = auto() + """ No scope; facts are never saved this is most often used with Schedule.ALWAYS """ + + class EnumerateModule(BaseModule): """Base class for all enumeration modules. @@ -96,7 +115,8 @@ class EnumerateModule(BaseModule): """ List of fact types which this module is capable of providing """ PLATFORM: typing.List[typing.Type[Platform]] = [] """ List of supported platforms for this module """ - + SCOPE: Scope = Scope.HOST + """ Defines the scope for this fact (either host or session) """ SCHEDULE: Schedule = Schedule.ONCE """ Determine the run schedule for this enumeration module """ @@ -122,6 +142,66 @@ class EnumerateModule(BaseModule): } """ Arguments accepted by all enumeration modules. This **should not** be overridden. """ + def _get_cached(self, session: "pwncat.manager.Session"): + """Retrieve the cached items for this module in the specified scope""" + + if self.SCOPE is Scope.HOST: + return [fact for fact in session.target.facts if fact.source == self.name] + + if self.SCOPE is Scope.SESSION: + return [fact for fact in session.facts if fact.source == self.name] + + return [] + + def _clear_cache(self, session: "pwncat.manager.Session"): + """Clear the cache based on the current scope""" + + if self.SCOPE is Scope.HOST: + session.target.facts = [ + fact for fact in session.target.facts if fact.source != self.name + ] + + if self.SCOPE is Scope.SESSION: + session.facts = [fact for fact in session.facts if fact.source != self.name] + + return [] + + def _mark_complete(self, session: "pwncat.manager.Session"): + """Mark this enumeration as complete for the scope and current schedule context""" + + if self.SCOPE is Scope.HOST: + state = session.target.enumerate_state + elif self.SCOPE is Scope.SESSION or self.SCOPE is Scope.NONE: + state = session.enumerate_state + + if self.SCHEDULE is Schedule.ONCE: + state[self.name] = True + elif self.SCHEDULE is Schedule.PER_USER: + if self.name not in state: + state[self.name] = [session.platform.getuid()] + elif session.platform.getuid() not in state[self.name]: + state[self.name].append(session.platform.getuid()) + + def _check_complete(self, session: "pwncat.manager.Session"): + """Check if this enumeration has already run for this scope and schedule context""" + + if self.SCHEDULE is Schedule.ALWAYS: + return False + + if self.SCOPE is Scope.HOST: + state = session.target.enumerate_state + elif self.SCOPE is Scope.SESSION or self.SCOPE is Scope.NONE: + state = session.enumerate_state + + if self.name not in state: + return False + elif self.SCHEDULE is Schedule.ONCE: + return True + elif self.SCHEDULE is Schedule.PER_USER: + return session.platform.getuid() in state[self.name] + + return False + def run( self, session: "pwncat.manager.Session", @@ -149,49 +229,34 @@ class EnumerateModule(BaseModule): target = session.target if clear: - # Filter out all facts which were generated by this module - target.facts = persistent.list.PersistentList( - (f for f in target.facts if f.source != self.name) - ) - - # Remove the enumeration state if available - if self.name in target.enumerate_state: - del target.enumerate_state[self.name] - + self._clear_cache(session) return # Yield all the know facts which have already been enumerated if cache and types: cached = [ f - for f in target.facts - if f.source == self.name - and any( + for f in self._get_cached(session) + if any( any(fnmatch.fnmatch(item_type, req_type) for req_type in types) for item_type in f.types ) ] elif cache: - cached = [f for f in target.facts if f.source == self.name] + cached = self._get_cached(session) else: cached = [] yield from cached # Check if the module is scheduled to run now - if (self.name in target.enumerate_state) and ( - (self.SCHEDULE == Schedule.ONCE and self.name in target.enumerate_state) - or ( - self.SCHEDULE == Schedule.PER_USER - and session.platform.getuid() in target.enumerate_state[self.name] - ) - ): + if self._check_complete(session): return for item in self.enumerate(session): # Allow non-fact status updates - if isinstance(item, Status) or self.SCHEDULE == Schedule.NOSAVE: + if isinstance(item, Status): yield item continue @@ -200,7 +265,7 @@ class EnumerateModule(BaseModule): if f == item: break else: - target.facts.append(item) + session.register_fact(item, self.SCOPE, commit=False) # Don't yield the actual fact if we didn't ask for this type if not types or any( @@ -215,13 +280,7 @@ class EnumerateModule(BaseModule): else: yield Status(item.title(session)) - # Update state for restricted modules - if self.SCHEDULE == Schedule.ONCE: - target.enumerate_state[self.name] = True - elif self.SCHEDULE == Schedule.PER_USER: - if self.name not in target.enumerate_state: - target.enumerate_state[self.name] = persistent.list.PersistentList() - target.enumerate_state[self.name].append(session.platform.getuid()) + self._mark_complete(session) def enumerate( self, session: "pwncat.manager.Session" diff --git a/pwncat/modules/windows/enumerate/system/processes.py b/pwncat/modules/windows/enumerate/system/processes.py index 93d47ea..ffad128 100644 --- a/pwncat/modules/windows/enumerate/system/processes.py +++ b/pwncat/modules/windows/enumerate/system/processes.py @@ -7,7 +7,7 @@ import rich.markup from pwncat.db import Fact from pwncat.modules import Status, ModuleFailed from pwncat.platform.windows import Windows, PowershellError -from pwncat.modules.enumerate import Schedule, EnumerateModule +from pwncat.modules.enumerate import Scope, Schedule, EnumerateModule class ProcessData(Fact): @@ -112,7 +112,8 @@ class Module(EnumerateModule): PROVIDES = ["system.processes"] PLATFORM = [Windows] # We don't save process results. They're volatile. Maybe this should be `Schedule.ALWAYS` anyway though? :shrug: - SCHEDULE = Schedule.NOSAVE + SCHEDULE = Schedule.ALWAYS + SCOPE = Scope.NONE def enumerate(self, session): diff --git a/pwncat/platform/windows.py b/pwncat/platform/windows.py index 45306a9..bcf8dc0 100644 --- a/pwncat/platform/windows.py +++ b/pwncat/platform/windows.py @@ -20,7 +20,6 @@ import stat import time import base64 import shutil -import signal import hashlib import pathlib import tarfile