Fix transform request handling from TDS

src/canari/commands/common.py

@@ -2,6 +2,8 @@
 
 from distutils.command.install import install
 from distutils.dist import Distribution
+from setuptools import find_packages
+from pkgutil import iter_modules
 from argparse import Action
 from datetime import datetime
 from string import Template
@@ -15,7 +17,7 @@
 
 from canari.commands.framework import Command
 from canari.config import CanariConfigParser
-
+from canari.utils.console import highlight
 
 __author__ = 'Nadeem Douba'
 __copyright__ = 'Copyright 2012, Canari Project'
@@ -69,7 +71,9 @@ def get_bin_dir():
     """
     Returns the absolute path of the installation directory for the Canari scripts.
     """
-    d = install(Distribution())
+    # re-import so we pass an isinstance check for Distribution
+    from distutils.dist import Distribution as MyDistribution
+    d = install(MyDistribution())
     d.finalize_options()
     return d.install_scripts
 
@@ -175,26 +179,112 @@ def project_root():
     raise ValueError('Unable to determine project root.')
 
 
-def project_tree():
-    root = project_root()
+def project_tree(package=None):
+    """Returns a dict of the project tree.
+
+    Will try and look for local/source packages first, and if it fails to find
+    a valid project root, it will look for system installed packages instead.
+
+    Returns a dictionary with the following fields:
+    - root: Path of the canari root folder or None if not applicable.
+    - src: Path of the folder containing the package.
+    - pkg: Path of the actual package.
+    - pkg_name: Name of the package, which details are returned about.
+    - resources: Path of the resources folder inside the package.
+    - transforms: Path of the transforms folder inside the package.
+    """
 
+    # Default values for the returned fields.
     tree = dict(
-        root=root,
-        # src is always directly under root
-        src=os.path.join(root, 'src'),
+        root=None,
+        src=None,
         pkg=None,
+        pkg_name=None,
         resources=None,
-        transforms=None
+        transforms=None,
     )
 
-    for base, dirs, files in os.walk(tree['src']):
-        if 'resources' in dirs:
-            tree['pkg'] = base
-        elif base.endswith('resources'):
-            tree['resources'] = base
-        elif base.endswith('transforms'):
-            tree['transforms'] = base
 
+    try:
+        root = project_root()
+
+        # TODO: The 'src' folder is currently harcoded inside setup.py. People
+        # may change this and thus we should probably read this value from
+        # '.canari', so the user may change this freely.
+
+        # Using find_packages we don't risk having to deal with the *.egg-info
+        # folder and trying to make a best guess at what folder is a actual
+        # source code, tests, or something else.
+        packages = filter(lambda pkg: pkg.find('.') < 0, find_packages('src'))
+        if package is None and len(packages) == 1:
+            # No package was specified by the user and there is only one
+            # possibility, so silently choose that one.
+            package = packages[0]
+        elif package not in packages:
+            # The supplied package was not found or not specified (None).  List
+            # the found packages and make the user choose the correct one.
+            if package is not None:
+                print "{warning} You specified a specific transform package, but " \
+                    "it does {_not_} exist inside this canari source directory. " \
+                    "\nPerhaps you ment to refer to an already installed package?\n" \
+                        .format(warning = highlight('[warning]', 'red', False),
+                                _not_= highlight('not', None, True))
+
+            print "The possible transform packages inside this canari root directory are:"
+            print 'Root dir: %s' % root
+            n = parse_int('Choose a package', packages, default=0)
+            package = packages[n]
+
+        #else: the user supplied package name is already a valid one, and the
+        #one the user picked.. so all is good.
+        assert package is not None, 'Fatal error: No package has been found or choosen!'
+
+        # Update the tree dict with all relevant information for this source package
+        tree['root'] = root
+        # Again 'src' is hardcooded in setup.py
+        tree['src'] = os.path.join(tree['root'], 'src')
+        tree['pkg'] = os.path.join(tree['src'], package)
+    except ValueError as ve:
+        # If we can't locate the project root, then we are not within a (source)
+        # canari project folder and thus we will try and look for installed
+        # packages instead.
+        for module_importer, name, ispkg in iter_modules():
+            # module_importer is most likely a pkgutils.ImpImporter instance (or
+            # the like) that has been initialised with a path that matches the
+            # (egg) install directory of the current module being iterated.
+            # Thus any calls to functions (e.g., find_module) on this instance
+            # will be limited to that path (i.e., you can't load arbitrary
+            # packages from it).
+            if name == package:
+                # Installed packages, don't have a (canari) 'root' folder.
+                # However it seems that (atleast) installed eggs have a form of
+                # 'src' folder named #pkg_name#-#pkg_version#-#py_version#.egg.
+                # This folder (generally) contains two folders: #pkg_name# and
+                # EGG-INFO
+                tree['src'] = module_importer.path
+                tree['pkg'] = module_importer.find_module(package).filename
+
+                break # No need to keep searching.
+
+        if tree['src'] is None:
+            # We didn't find the user supplied package name in the list of
+            # installed packages.
+            raise ValueError("You are not inside a canari root directory ('%s'), "
+                             "and it was not possible to locate " "the given package "
+                             "'%s' among the list of installed packages."
+                             % (os.getcwd(), package))
+
+
+    tree['pkg_name'] = package
+    # A transform packages structure is expected to have a 'pkg_name.resources'
+    # and 'pkg_name.transforms', thus we won't dynamically look for these as
+    # everything else will break, if they can't be imported as such.
+
+    # TODO: Here be dragons. Does python3 module madness break this assumption
+    # with its new fancy features of ways to have modules not nessesarily
+    # stricly tied to the file system?
+    tree['resources'] = os.path.join(tree['pkg'], 'resources')
+    tree['transforms'] = os.path.join(tree['pkg'], 'transforms')
     return tree
 
 
@@ -244,4 +334,4 @@ def __enter__(self):
         return self
 
     def __exit__(self, type_, value, tb):
-        os.chdir(self.original_dir)
+        os.chdir(self.original_dir)

src/canari/commands/list_commands.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 
 from common import canari_main
-from canari.maltego.utils import highlight
+from canari.utils.console import highlight
 from framework import SubCommand
 
 __author__ = 'Nadeem Douba'
@@ -25,4 +25,4 @@ def list_commands(opts):
     k = cmds.keys()
     k.sort()
     for i in k:
-        print ('%s - %s' % (highlight(i, 'green', True), cmds[i].description))
+        print ('%s - %s' % (highlight(i, 'green', True), cmds[i].description))

src/canari/commands/list_transforms.py

@@ -1,16 +1,16 @@
 #!/usr/bin/env python
 
 import os
-from canari.maltego.utils import highlight
+from canari.utils.console import highlight
 from canari.pkgutils.transform import TransformDistribution
 
-from common import (canari_main, uproot, pushd)
+from common import (canari_main, pushd, project_tree)
 from framework import SubCommand, Argument
 
 
 __author__ = 'Nadeem Douba'
 __copyright__ = 'Copyright 2012, Canari Project'
-__credits__ = []
+__credits__ = ['Jesper Reenberg']
 
 __license__ = 'GPL'
 __version__ = '0.6'
@@ -19,36 +19,77 @@
 __status__ = 'Development'
 
 
-# Extra sauce to parse args
 def parse_args(args):
+    args.ptree = project_tree(package = args.package)
+    args.package = args.ptree['pkg_name']
+    # We specifically don't update 'args' with any of the info from ptree.  This
+    # way we always know exactly what information was specified by the user.
     return args
 
-
 # Argument parser
 @SubCommand(
     canari_main,
-    help="Installs and configures canari transform packages in Maltego's UI",
-    description="Installs and configures canari transform packages in Maltego's UI"
+    help="List transforms inside the given transform package",
+    description="List transforms inside a given transform package (<package>).  "
+        "Python 'import' ordering is used, thus a specified directory (--dir) "
+        "will supersede the current working directory which superseeds installed "
+        "packages, as long as a canari project is found in any of the two.  If "
+        "no package name is specified, then all possible transform packages "
+        "inside the found canari project is listed."
 )
 @Argument(
     'package',
     metavar='<package>',
-    help='the name of the canari transforms package to install.'
+    nargs='?',
+    default=None,
+    help="the name of the canari transform package to list transforms from.  If"
+        "no canari project is located, then the installed modules is searched."
 )
 @Argument(
-    '-w',
-    '--working-dir',
-    metavar='[working dir]',
-    default=None,
-    help="the path that will be used as the working directory for "
-         "the transforms being installed (default: ~/.canari/)"
+    '-d',
+    '--dir',
+    metavar='[dir]',
+    default=os.getcwd(),
+    help="if supplied, the path will owerwrite the current working directory when "
+         "searching for canari projects."
 )
+
+
 def list_transforms(args):
 
-    opts = parse_args(args)
+    # TODO: project_tree may raise an exception if either project_root can't be
+    # determined or if we can't find the package as an installed package.
+    # Atleast the create-transform command calls this function without handling
+    # the possible exception.  What is the best sollution?
+
+    # TODO: create-transform takes an argument --transform-dir which can be used
+    # to control where to place the transform template. This breaks the new
+    # assumption of the 'transforms' folder always being inside the 'pkg'
+    # folder.  However this is an assumption all over the place, so this
+    # parameter doesn't really make much sense?
+
+    # TODO: There are most likely many commands with similar problems
+    # (above). and perhaps they should be updated to use the below template and
+    # have their argument updated to -d/--dir instead with CWD as the default
+    # value.
+
+    # TODO: Perhaps we should introduce a 'create' command that will just make
+    # an empty canari root dir (project). Inside this we can then call
+    # create-package a number of times to generate all the desired
+    # packages. This could even be automated for N-times during the call to
+    # 'create'. 'create-package' can even still default to call 'create' if not
+    # inside a canari root directory, to preserve backwards compatability.
+
+    # TODO: Handle hyphening of package names. When creating them and when
+    # trying to access them.  This goes for project_tree, it should change '-'
+    # with '_' in the package name.
 
     try:
-        with pushd(opts.working_dir or os.getcwd()):
+        with pushd(args.dir):
+            opts = parse_args(args)
+
+        with pushd(args.ptree['src']):
+
             transform_package = TransformDistribution(opts.package)
             for t in transform_package.transforms:
                 print ('`- %s: %s' % (highlight(t.__name__, 'green', True), t.dotransform.description))
@@ -62,4 +103,4 @@ def list_transforms(args):
                 print ''
     except ValueError, e:
         print str(e)
-        exit(-1)
+        exit(-1)

src/canari/commands/run_server.py

@@ -48,7 +48,9 @@ def message(m, response):
             m.replace(url, new_url, 1)
         v = m
     else:
-        v = MaltegoMessage(m).render(fragment=True)
+        mm = MaltegoMessage()
+        mm.message = m
+        v = mm.render(fragment=True)
         # Get rid of those nasty unicode 32 characters
     response.wfile.write(v)
 
@@ -80,7 +82,7 @@ def dotransform(self, transform, valid_input_entity_types):
                 return
             request_str = self.rfile.read(int(self.headers['Content-Length']))
 
-            msg = MaltegoTransformRequestMessage.parse(request_str).message
+            msg = MaltegoMessage.parse(request_str).message
 
             e = msg.entity
             entity_type = e.type

src/canari/commands/shell.py

@@ -9,7 +9,8 @@
 from common import canari_main, fix_pypath, fix_binpath, import_package, pushd
 from framework import SubCommand, Argument
 from canari.config import config
-from canari.maltego.utils import highlight, console_message, local_transform_runner
+from canari.maltego.utils import console_message, local_transform_runner
+from canari.utils.console import highlight
 import canari
 
 

src/canari/config.py

@@ -95,7 +95,7 @@ def __setitem__(self, key, value):
         section, option = key.split('/', 1)
         if not self.has_section(section):
             self.add_section(section)
-        self.set(section, option, value)
+        self.set(section, option, value.value)
 
 
 config = CanariConfigParser()
@@ -104,4 +104,4 @@ def __setitem__(self, key, value):
 lconf = path.join(getcwd(), 'canari.conf')
 
 config.read([dconf, lconf])
-config.read(config['default/configs'])
+config.read(config['default/configs'])