diff --git a/pupy/modules/psh.py b/pupy/modules/psh.py new file mode 100644 index 00000000..c65380b4 --- /dev/null +++ b/pupy/modules/psh.py @@ -0,0 +1,169 @@ +# -*- coding: utf-8 -*- + +from pupylib.PupyModule import * +from pupylib.utils.term import consize +from argparse import REMAINDER +from pupylib.utils.rpyc_utils import obtain +from pupylib.PupyCmd import PupyCmd +from os import path +from rpyc import GenericException + +__class_name__ = 'PowershellManager' + +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) + +@config(compat='windows', category='admin') +class PowershellManager(PupyModule): + ''' Load/Execute Powershell scripts ''' + + dependencies = { + 'windows': [ 'powershell' ] + } + + def init_argparse(self): + self.arg_parser = PupyArgumentParser( + prog='psh', description=self.__doc__ + ) + + width, _ = consize() + + commands = self.arg_parser.add_subparsers(title='actions') + loaded = commands.add_parser('loaded', help='List preapred powershell contexts') + loaded.add_argument('context', nargs='?', help='Check is context with specified name loaded') + loaded.set_defaults(name='loaded') + + load = commands.add_parser('load', help='Create new powershell context or load more scripts into existing') + load.add_argument('-F', '--force', action='store_true', default=False, help='Destroy old context if exists') + load.add_argument('-64', '--try-64', action='store_true', default=False, help='Try amd64 if possible') + load.add_argument('-2', '--try-v2', action='store_true', default=None, help='Try version 2 if possible') + load.add_argument('-W', '--width', default=width, type=int, help='Set output line width') + load.add_argument('-D', '--daemon', action='store_true', default=False, help='Start in "daemon" mode') + load.add_argument('context', help='Context name') + load.add_argument('source', help='Path to PS1 script (local to pupy)') + load.set_defaults(name='load') + + iex = commands.add_parser('iex', help='Invoke expression in context') + iex.add_argument('-T', '--timeout', default=None, type=float, help='Set timeout for result retrieval') + iex.add_argument('-B', '--background', default=False, action='store_true', + help='Evaluate in background (async)') + iex.add_argument('context', help='Context name') + iex.add_argument('expression', nargs=REMAINDER, help='Expression to evaluate') + iex.set_defaults(name='iex') + + unload = commands.add_parser('unload', help='Destroy context') + unload.add_argument('context', help='Context name') + unload.set_defaults(name='unload') + + result = commands.add_parser('result', help='Retrieve result by request id from context') + result.add_argument('context', help='Context name') + result.add_argument('rid', type=int, help='Request id') + result.set_defaults(name='result') + + results = commands.add_parser('results', help='Retrieve ready RIDs from all contexts') + results.set_defaults(name='results') + + killall = commands.add_parser('killall', help='Destroy all powershell contexts') + killall.set_defaults(name='killall') + + def run(self, args): + powershell = self.client.conn.modules['powershell'] + + if args.name == 'loaded': + if args.context: + if powershell.loaded(args.context): + self.success('{} is loaded'.format(args.context)) + else: + self.error('{} is not loaded'.format(args.context)) + else: + contexts = obtain(powershell.loaded()) + for context in contexts: + self.success('{}'.format(context)) + + elif args.name == 'load': + script = path.expandvars(path.expanduser(args.source)) + if not path.exists(script): + self.error('Script file not found: {}'.format(script)) + return + + with open(script) as input: + content = input.read() + + try: + powershell.load( + args.context, content, args.force, args.try_64, + args.daemon, args.width, args.try_v2 + ) + except Exception, e: + self.error('load: {}'.format(e)) + + elif args.name == 'iex': + expression = ' '.join(args.expression) + + if not powershell.loaded(args.context): + self.error('Context {} is not loaded'.format(args.context)) + return + + if not expression: + self.warning('Empty expression') + return + + try: + result = powershell.call( + args.context, expression, timeout=args.timeout, async=args.background + ) + + if args.background: + self.warning('Queued: Context: {} RID: {}'.format(args.context, result.rid)) + else: + output, rest = result + + if rest: + self.warning(rest) + if output: + self.log(output) + + except GenericException as e: + if type(e).__name__ == 'powershell.PowershellTimeout': + self.error('iex: timeout: Context: {} RID: {}'.format(args.context, e.args[0])) + else: + self.error('iex: {}'.format(e)) + + except Exception, e: + self.error('iex: {}'.format(e)) + + elif args.name == 'unload': + try: + powershell.unload(args.context) + except Exception, e: + self.error('unload: {}'.format(e)) + + elif args.name == 'result': + result = powershell.result(args.context, args.rid) + if not result: + self.error('Result {} does not exists in {}'.format(args.rid, args.context)) + else: + output, rest = result + if rest: + self.warning(rest) + if output: + self.log(output) + + elif args.name == 'results': + if not obtain(powershell.loaded()): + self.error('No scripts loaded') + return + + results = obtain(powershell.results.fget()) + objects = [ + { + 'CONTEXT': ctx, + 'RIDS': ', '.join([str(x) for x in rids]) + } for ctx, rids in results.iteritems() + ] + self.stdout.write( + PupyCmd.table_format(objects, wl=[ + 'CONTEXT', 'RIDS' + ])) + + elif args.name == 'killall': + powershell.stop() diff --git a/pupy/packages/windows/all/powershell.py b/pupy/packages/windows/all/powershell.py index c1eef547..866d7395 100644 --- a/pupy/packages/windows/all/powershell.py +++ b/pupy/packages/windows/all/powershell.py @@ -15,13 +15,19 @@ import codecs class PowerHostUninitialized(Exception): pass -class PowerShellUninitialized(Exception): +class PowershellUninitialized(Exception): pass -class PowerShellInitializationFailed(Exception): +class PowershellInitializationFailed(Exception): pass -class PowerShellV2NotInstalled(Exception): +class PowershellContextUnregistered(Exception): + pass + +class PowershellV2NotInstalled(Exception): + pass + +class PowershellTimeout(Exception): pass class Request(object): @@ -36,6 +42,10 @@ class Request(object): self._rid = rid self._completed = False + @property + def rid(self): + return self._rid + @property def expression(self): return self._expression @@ -46,7 +56,9 @@ class Request(object): @property def result(self): - self._event.wait(timeout=self._timeout) + if not self._event.wait(timeout=self._timeout): + raise PowershellTimeout(self._rid) + result = self._result if self._rid in self._storage: @@ -88,12 +100,11 @@ class Request(object): self._event.set() -class PowerShell(threading.Thread): +class Powershell(threading.Thread): def __init__(self, host, name, content, try_x64=False, daemon=False, width=None, v2=True): - super(PowerShell, self).__init__() + super(Powershell, self).__init__() self.daemon = True - self._content = zlib.compress(content) if content else None self._try_x64 = try_x64 self._completed = threading.Event() self._initialized = False @@ -109,7 +120,7 @@ class PowerShell(threading.Thread): self._v2 = v2 if try_x64: - native = ur'C:\Windows\SysNative\WindowsPowerShell\v1.0\powershell.exe' + native = ur'C:\Windows\SysNative\WindowsPowershell\v1.0\powershell.exe' if os.path.exists(native): self._executable = native @@ -117,13 +128,13 @@ class PowerShell(threading.Thread): self._executable, u'-W', u'hidden', u'-I', u'Text', u'-C', u'-' ] - self._initialize() + self._initialize(content) @property def v2(self): return self._v2 - def _initialize(self): + def _initialize(self, content): if self._pipe: return @@ -155,26 +166,26 @@ class PowerShell(threading.Thread): if 'Version v2.0.50727 of the .NET Framework is not installed'.encode('UTF-16LE') in data: self.stop() - raise PowerShellV2NotInstalled() + raise PowershellV2NotInstalled() elif not data or not preamble_complete in data: print "First line: ", repr(data) print '.NET Framework is not installed' in data self.stop() - raise PowerShellInitializationFailed() + raise PowershellInitializationFailed() - if not self._content: + if not content: return - content = zlib.decompress(self._content) + self.load(content) + + def load(self, content): content = re.sub('Write-Host ', 'Write-Output ', content, flags=re.I) - self._invoke_expression(content) - def _invoke_expression(self, content, dest=None, pipe=None): if not self._pipe: - raise PowerShellUninitialized() + raise PowershellUninitialized() if not content: return @@ -273,7 +284,7 @@ class PowerShell(threading.Thread): sol_at+len(SOL):-(len(EOL)+1) ].strip(), response[:sol_at] - def execute(self, expression, async=False, timeout=None): + def execute(self, expression, async=False, timeout=None, wait=None): if self._daemon: async = True @@ -281,6 +292,12 @@ class PowerShell(threading.Thread): self._execute(expression) return self._daemon_request + if timeout: + if wait is None: + wait = not async + + async = True + if self._queue: request = Request( self._rid, @@ -295,7 +312,8 @@ class PowerShell(threading.Thread): self._rid += 1 self._queue.put(request) - if async: + + if async and not wait: return request else: return request.result @@ -304,7 +322,7 @@ class PowerShell(threading.Thread): self._queue = Queue.Queue() self._host.results[self._name] = {} self.start() - return self.execute(expression, async, timeout) + return self.execute(expression, async, timeout, wait) else: return self._execute(expression) @@ -358,18 +376,27 @@ class PowerHost(object): self._powershells[name].stop() try: - self._powershells[name] = PowerShell( + self._powershells[name] = Powershell( self, name, content, try_x64, daemon, width, v2 ) - except PowerShellV2NotInstalled: + except PowershellV2NotInstalled: self._v2 = False - self._powershells[name] = PowerShell( + self._powershells[name] = Powershell( self, name, content, try_x64, daemon, width, False ) - def registered(self, name): - return name in self._powershells + def load(self, name, content): + if not name in self._powershells: + raise PowershellUninitialized() + + self._powershells[name].load(content) + + def registered(self, name=None): + if name: + return name in self._powershells + else: + return self._powershells.keys() def unregister(self, name): if not name in self._powershells: @@ -417,10 +444,13 @@ class PowerHost(object): def stopped(self): return not self._powershells -def loaded(name): +def loaded(name=None): powershell = pupy.manager.get(PowerHost) if not powershell: - return False + if not name: + return [] + else: + return False return powershell.registered(name) @@ -431,7 +461,10 @@ def load(name, content, force=False, try_x64=False, daemon=False, width=None, v2 if not powershell: raise PowerHostUninitialized() - powershell.register(name, content, force, try_x64, daemon, width, v2) + if loaded(name) and not force: + powershell.load(name, content) + else: + powershell.register(name, content, force, try_x64, daemon, width, v2) def unload(name): powershell = pupy.manager.get(PowerHost) @@ -457,13 +490,34 @@ def call(name, expression, async=False, timeout=None, content=None, try_x64=Fals if content: unload(name) +def result(name, rid): + rid = int(rid) + + powershell = pupy.manager.get(PowerHost) + if not powershell: + raise PowerHostUninitialized() + + if not loaded(name): + PowershellContextUnregistered() + + results = powershell.results + if not name in results or not rid in results[name]: + return None + + result = results[name][rid] + del results[name][rid] + + return result + @property def results(): powershell = pupy.manager.get(PowerHost) if not powershell: raise PowerHostUninitialized() - return powershell.results + return { + ctx:results.keys() for ctx, results in powershell.results.iteritems() + } def stop(): pupy.manager.stop(PowerHost)