1
0
mirror of https://github.com/calebstewart/pwncat.git synced 2024-11-24 01:25:37 +01:00

Added enumeration scopes

This commit is contained in:
Caleb Stewart 2021-06-18 22:04:12 -04:00
parent 270f6793ad
commit 3e501d2957
7 changed files with 118 additions and 84 deletions

View File

@ -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

View File

@ -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):

View File

@ -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)
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"""

View File

@ -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")

View File

@ -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"

View File

@ -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):

View File

@ -20,7 +20,6 @@ import stat
import time
import base64
import shutil
import signal
import hashlib
import pathlib
import tarfile