From 7503d1caa5e11c67d88f97fb2a13b5803cca74f5 Mon Sep 17 00:00:00 2001 From: Nadeem Douba Date: Thu, 6 Dec 2012 23:16:16 -0500 Subject: [PATCH] Added generate-entities command that reads from both mtz and Entity directory --- src/canari/commands/__init__.py | 3 +- src/canari/commands/generate_entities.py | 207 +++++++++++++++++++++++ 2 files changed, 209 insertions(+), 1 deletion(-) create mode 100644 src/canari/commands/generate_entities.py diff --git a/src/canari/commands/__init__.py b/src/canari/commands/__init__.py index 92243c4..19ea954 100644 --- a/src/canari/commands/__init__.py +++ b/src/canari/commands/__init__.py @@ -25,5 +25,6 @@ __all__ = [ 'delete_transform', 'rename_transform', 'list_commands', - 'run_server' + 'run_server', + 'generate_entities' ] \ No newline at end of file diff --git a/src/canari/commands/generate_entities.py b/src/canari/commands/generate_entities.py new file mode 100644 index 0000000..58cc839 --- /dev/null +++ b/src/canari/commands/generate_entities.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python + +from common import detect_settings_dir, cmd_name +from ..maltego.entities import Entity + +from xml.etree.cElementTree import XML +from argparse import ArgumentParser +from os import walk, path as ospath +from zipfile import ZipFile +from re import sub, findall +from imp import load_source + + +__author__ = 'Nadeem Douba' +__copyright__ = 'Copyright 2012, Canari Project' +__credits__ = ['Nadeem Douba'] + +__license__ = 'GPL' +__version__ = '0.1' +__maintainer__ = 'Nadeem Douba' +__email__ = 'ndouba@gmail.com' +__status__ = 'Development' + + + +__author__ = 'Nadeem Douba' +__copyright__ = 'Copyright 2012, Canari Project' +__credits__ = [] + +__license__ = 'GPL' +__version__ = '0.1' +__maintainer__ = 'Nadeem Douba' +__email__ = 'ndouba@gmail.com' +__status__ = 'Development' + + +parser = ArgumentParser( + description='Converts Maltego entity definition files to Canari python classes. Excludes Maltego built-in entities.', + usage='canari %s [options]' % cmd_name(__name__) +) + + +parser.add_argument( + 'outfile', + metavar='', + help='Which file to write the output to.', + default='entities.py', + nargs='?' +) + +parser.add_argument( + '--mtz-file', + '-m', + metavar='', + help='A *.mtz file containing an export of Maltego entities.', + required=False +) + +parser.add_argument( + '--exclude-namespace', + '-e', + metavar='', + help='Name of Maltego entity namespace to ignore. Can be defined multiple times.', + required=False, + action='append', + default=['maltego', 'maltego.affiliation'] +) + +parser.add_argument( + '--namespace', + '-n', + metavar='', + help='Name of Maltego entity namespace to generate entity classes for. Can be defined multiple times.', + required=False, + action='append', + default=[] +) + +parser.add_argument( + '--append', + '-a', + help='Whether or not to append to the existing *.py file.', + action='store_true', + default=False +) + +parser.add_argument( + '--entity', + '-E', + metavar='', + help='Name of Maltego entity to generate Canari python class for.', + required=False, + action='append', + default=[] +) + + +def help(): + parser.print_help() + + +def description(): + return parser.description + + +def parse_args(args): + return parser.parse_args(args) + + +def normalize_fn(fn): + # Get rid of starting underscores or numbers and bad chars for var names in python + return sub(r'[^A-Za-z0-9]', '', sub(r'^[^A-Za-z]+', '', fn)) + + +def diff(fn): + m = load_source('entities', fn) + l = [] + for c in dir(m): + try: + i = m.__dict__[c]('load') + if isinstance(i, Entity): + l.append(i) + except TypeError: + pass + return l + + +class DirFile(object): + + def __init__(self, path): + self.path = path + + def namelist(self): + l = [] + for base, dirs, files in walk(self.path): + l.extend([ ospath.join(base, f) for f in files ]) + return l + + def open(self, fname): + return file(fname) + + +def run(args): + + opts = parse_args(args) + + ar = DirFile( + ospath.join(detect_settings_dir(), 'config', 'Maltego', 'Entities') + ) if opts.mtz_file is None else ZipFile(opts.mtz_file) + + entities = filter(lambda x: x.endswith('.entity'), ar.namelist()) + + nses = dict() + + el = [] + if opts.append: + l = diff(opts.outfile) + el.extend([i.type for i in l]) + for i in l: + if i.type.endswith('Entity'): + nses[i.namespace] = i.__class__.__name__ + + fd = open(opts.outfile, 'ab' if opts.append else 'wb') + + if opts.append: + fd.write('\n\n') + else: + fd.write('#!/usr/bin/env python\n\nfrom canari.maltego.entities import EntityField, Entity\n\n\n') + + for e in entities: + xml = XML(ar.open(e).read()) + id_ = xml.get('id') + + if (opts.entity and id_ not in opts.entity) or id_ in el: + continue + + ens = id_.split('.') + + base_classname = None + namespace = '.'.join(ens[:-1]) + name = ens[-1] + classname = name + + if (opts.namespace and namespace not in opts.namespace) or namespace in opts.exclude_namespace: + continue + + if namespace not in nses: + base_classname = '%sEntity' % (''.join([ n.title() for n in ens[:-1] ])) + nses[namespace] = base_classname + + fd.write('class %s(Entity):\n namespace = %s\n\n' % (base_classname, repr(namespace))) + else: + base_classname = nses[namespace] + + + for f in xml.findall('Properties/Fields/Field'): + fields = [ + 'name=%s' % repr(f.get('name')), + 'propname=%s' % repr(normalize_fn(f.get('name'))), + 'displayname=%s' % repr(f.get('displayName')) + + ] + fd.write('@EntityField(%s)\n' % ', '.join(fields)) + + fd.write('class %s(%s):\n pass\n\n\n' % (classname, base_classname)) + + fd.close() \ No newline at end of file -- 2.45.1