From b0f8d56ef0b2b63faed36f43f3677044cf3f3051 Mon Sep 17 00:00:00 2001
From: allfro <ndouba@gmail.com>
Date: Sun, 21 Oct 2012 14:19:16 -0400
Subject: [PATCH] Fixed canari command shell and debug to include sudo
 functionality.

---
 src/canari/commands/common.py          |  5 +++
 src/canari/commands/debug_transform.py |  6 +--
 src/canari/commands/shell.py           | 57 ++++++++++++++------------
 3 files changed, 38 insertions(+), 30 deletions(-)

diff --git a/src/canari/commands/common.py b/src/canari/commands/common.py
index 7e19c3d..3ffd126 100644
--- a/src/canari/commands/common.py
+++ b/src/canari/commands/common.py
@@ -124,6 +124,11 @@ def import_transform(script):
     return __import__(script, globals(), locals(), ['dotransform'])
 
 
+def import_package(package):
+    fix_pypath()
+    return __import__(package, globals(), locals(), ['*'])
+
+
 def cmd_name(name):
     return name.replace('canari.commands.', '').replace('_', '-')
 
diff --git a/src/canari/commands/debug_transform.py b/src/canari/commands/debug_transform.py
index 34b02ae..f20fdca 100644
--- a/src/canari/commands/debug_transform.py
+++ b/src/canari/commands/debug_transform.py
@@ -5,8 +5,7 @@ from common import croak, import_transform, cmd_name, console_message, fix_binpa
 from ..maltego.utils import onterminate, parseargs
 from ..config import config
 
-from distutils.sysconfig import get_config_var
-from os import execvp, geteuid, name, path
+from os import execvp, geteuid, name
 from argparse import ArgumentParser
 from traceback import format_exc
 from sys import argv
@@ -66,13 +65,12 @@ def run(args):
     [transform, params, value, fields] = parseargs(['canari %s' % cmd_name(__name__)] + args)
 
     m = None
-    pysudo = path.join(get_config_var('BINDIR'), 'pysudo')
     fix_binpath(config['default/path'])
     try:
         m = import_transform(transform)
 
         if name == 'posix' and hasattr(m.dotransform, 'privileged') and geteuid():
-            execvp(pysudo, [pysudo] + list(argv))
+            execvp('sudo', ['sudo'] + list(argv))
             exit(-1)
 
         if hasattr(m, 'onterminate'):
diff --git a/src/canari/commands/shell.py b/src/canari/commands/shell.py
index f700f3e..355f259 100644
--- a/src/canari/commands/shell.py
+++ b/src/canari/commands/shell.py
@@ -1,14 +1,15 @@
 #!/usr/bin/env python
 
-from common import console_message, cmd_name, highlight, fix_pypath, fix_binpath
+from common import console_message, cmd_name, highlight, fix_pypath, fix_binpath, import_package
 from ..maltego.message import MaltegoTransformResponseMessage
 from ..config import config
 
+from distutils.sysconfig import get_config_var
+from os import path, name, geteuid, execvp
 from code import InteractiveConsole
 from argparse import ArgumentParser
-from os import path, environ
 from atexit import register
-from re import sub, match
+from sys import argv
 import readline
 
 
@@ -43,37 +44,41 @@ def description():
     return parser.description
 
 
-class MtgConsole(InteractiveConsole):
+class ShellCommand(object):
 
-    def __init__(self, package):
-        m = __import__(package, globals(), locals(), ['*'])
-        m.__dict__[MaltegoTransformResponseMessage.__name__] = MaltegoTransformResponseMessage
-        m.__dict__['message'] = console_message
-        InteractiveConsole.__init__(self, locals=m.__dict__)
-        self.init_history(path.expanduser('~/.mtgsh_history'))
+    def __init__(self, mod):
+        self.mod = mod
+        self.sudoargs = ['sudo'] + list(argv)
 
-    def raw_input(self, prompt):
-        line = InteractiveConsole.raw_input(self, prompt=highlight('mtg> ', None, True))
-        r = match(r'([^(]+)\((.*?)\)', line)
-        if r is not None:
-            g = r.groups()
-            if g[0] in self.locals:
-                line = eval(sub(r'([^\(]+)\((.*)\)', r'self.docall("\1", \2)', line))
-        return line
-
-    def docall(self, *args, **kwargs):
-        return """message(%s.dotransform(
+    def __call__(self, value, *args, **kwargs):
+        if name == 'posix' and hasattr(self.mod.dotransform, 'privileged') and geteuid():
+            print highlight("Need to be root to run this transform... sudo'ing...", 'green', True)
+            execvp('sudo', self.sudoargs)
+            return
+        return console_message(self.mod.dotransform(
             type(
                 'MaltegoTransformRequestMessage',
                 (object,),
-                {
-                    'value' : %s,
-                    'fields' : %s,
-                    'params' : %s
+                    {
+                    'value' : value,
+                    'params' : list(args),
+                    'fields' : kwargs
                 }
             )(),
             MaltegoTransformResponseMessage()
-        ))""" % (args[0], repr(args[1]) if len(args) != 1 else repr(''), repr(kwargs), repr(args[2:]))
+        ))
+
+
+class MtgConsole(InteractiveConsole):
+
+    def __init__(self, package):
+        package = import_package(package)
+        transforms = dict(dir=dir)
+        for name, mod in package.__dict__.iteritems():
+            if getattr(mod, 'dotransform', ''):
+                transforms[name] = ShellCommand(mod)
+        InteractiveConsole.__init__(self, locals=transforms)
+        self.init_history(path.expanduser('~/.mtgsh_history'))
 
     def init_history(self, histfile):
         readline.parse_and_bind('tab: complete')
-- 
2.45.1