diff --git a/.gitignore b/.gitignore index 55850b7..9fbace3 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ __pycache__/*g *~ .DS_Store +*.egg-info diff --git a/polymorph/UI/interface.py b/polymorph/UI/interface.py index bb1506a..88c4405 100644 --- a/polymorph/UI/interface.py +++ b/polymorph/UI/interface.py @@ -11,7 +11,7 @@ from polymorph.UI.command_parser import CommandParser from collections import OrderedDict import os -from os.path import dirname +# from os.path import dirname import polymorph @@ -21,7 +21,7 @@ class Interface(object): def __init__(self): self._poisoner = None - self._polym_path = dirname(polymorph.__file__) + self._polym_path = polymorph.settings.path @staticmethod def print_help(hdict): diff --git a/polymorph/UI/templateinterface.py b/polymorph/UI/templateinterface.py index fb74874..40d91b7 100644 --- a/polymorph/UI/templateinterface.py +++ b/polymorph/UI/templateinterface.py @@ -21,7 +21,7 @@ import polymorph.conditions import platform from shutil import copyfile - +from polymorph import settings, utils class TemplateInterface(Interface): """This class is responsible for parsing and respond to user commands in @@ -46,7 +46,7 @@ def __init__(self, template, index, poisoner=None): self._t = template self._index = index self._poisoner = poisoner - self._conds_path = dirname(polymorph.conditions.__file__) + self._conds_path = settings.paths['conditions'] def run(self): """Runs the interface and waits for user input commands.""" @@ -256,10 +256,10 @@ def _conditions(self, command, cond): # Adds a condition elif args['-a']: # Create the new file if not exists - if not os.path.isfile(join(self._conds_path, cond, args["-a"] + ".py")): + if not os.path.isfile(join(self._conds_path, settings.paths[cond], args["-a"] + ".py")): self._create_cond(cond, args["-a"]) ret = os.system("%s %s.py" % ( - args["-e"], join(self._conds_path, cond, args["-a"]))) + args["-e"], join(self._conds_path, settings.paths[cond], args["-a"]))) if ret != 0: Interface._print_error( "The editor is not installed or is not in the PATH") @@ -288,6 +288,7 @@ def _conditions(self, command, cond): name = file if name[-3:] == ".py": name = name[:-3] + self._add_cond(cond, name) try: self._add_cond(cond, name) Interface._print_info("Condition %s imported" % args['-i']) @@ -300,14 +301,12 @@ def _conditions(self, command, cond): def _create_cond(self, cond, name): """Creates a new condition with the initial source code.""" code = "def %s(packet):\n\n # your code here\n\n # If the condition is meet\n return packet" % name - with open("%s.py" % join(self._conds_path, cond, name), 'w') as f: + with open("%s.py" % join(settings.paths[cond], name), 'w') as f: f.write(code) def _add_cond(self, cond, name): """Adds a new condition to the `Template`.""" - m = importlib.import_module( - "polymorph.conditions.%s.%s" % (cond, name)) - importlib.reload(m) + m = utils.import_file(join(settings.paths[cond], "{}.py".format(name)), "polymorph.conditions.%s.%s" % (cond, name)) self._t.add_function(cond, name, getattr(m, dir(m)[-1])) @staticmethod diff --git a/polymorph/UI/tlistinterface.py b/polymorph/UI/tlistinterface.py index 919e193..821403d 100644 --- a/polymorph/UI/tlistinterface.py +++ b/polymorph/UI/tlistinterface.py @@ -14,7 +14,7 @@ from polymorph.deps.prompt_toolkit.auto_suggest import AutoSuggestFromHistory from polymorph.deps.prompt_toolkit.shortcuts import CompleteStyle import polymorph -from os.path import dirname, join +from os.path import join class TListInterface(Interface): @@ -37,7 +37,7 @@ def __init__(self, tlist, poisoner=None): # Class Attributes self._t = tlist self._poisoner = poisoner - self._polym_path = dirname(polymorph.__file__) + self._polym_path = polymorph.settings.path def run(self): """Runs the interface and waits for user input commands.""" diff --git a/polymorph/__init__.py b/polymorph/__init__.py index 8b13789..1e254f8 100644 --- a/polymorph/__init__.py +++ b/polymorph/__init__.py @@ -1 +1,4 @@ +from .settings import Settings +# PoC to create initial settings manager +settings = Settings() diff --git a/polymorph/settings.py b/polymorph/settings.py new file mode 100644 index 0000000..570f87d --- /dev/null +++ b/polymorph/settings.py @@ -0,0 +1,21 @@ +import os + +class Settings (object): + + def __init__(self): + self._poisoner = None + self.path = os.path.expanduser("~/.polymorph") + + self.paths = { + "": self.path, + "conditions": "{}/conditions".format(self.path), + "preconditions": "{}/conditions/preconditions".format(self.path), + "postconditions": "{}/conditions/postconditions".format(self.path), + "executions": "{}/conditions/executions".format(self.path), + "templates": "{}/templates".format(self.path), + } + + # TODO; Ensure that all App reuse this iface._polym_path instead of redefine it + for _path in self.paths.values(): + if not os.path.exists(_path): + os.makedirs(_path) diff --git a/polymorph/template.py b/polymorph/template.py index fb73370..f8e40bd 100644 --- a/polymorph/template.py +++ b/polymorph/template.py @@ -16,7 +16,7 @@ import polymorph.conditions from os import listdir from os.path import isfile, join - +from polymorph import settings class Template: """Main class that represents a template""" @@ -47,7 +47,7 @@ def __init__(self, name=None, raw=None, from_path=None, if from_path: self.read(from_path) # Path to conditions (precs, execs, posts) - self._conds_path = dirname(polymorph.conditions.__file__) + self._conds_path = dirname(settings.paths['conditions']) def __repr__(self): return "" % "/".join(self.layernames()) @@ -67,7 +67,13 @@ def write(self, path=None): """ if not path: path = "../templates/" + self._name.replace("/", "_") + ".json" - with open(path, 'w') as outfile: + + if not path.endswith(".json"): + path += ".json" + + template_file = "{}/{}".format(settings.paths['templates'], path) + + with open(template_file, 'w') as outfile: json.dump(self.dict(), outfile, indent=4) @property @@ -213,7 +219,7 @@ def del_execution(self, name): """ self.del_function('executions', name) - def get_function_source(self, func, name): + def get_function_source(self, cond, name): """Returns a precondition/postcondition/execution source code. Parameters @@ -229,7 +235,7 @@ def get_function_source(self, func, name): Source code of the function. """ - path = "%s/%s/%s.py" % (self._conds_path, func, name) + path = join(settings.paths[cond], "{}.py".format(name)) if os.path.isfile(path): return open(path).read() return "[!] File is not in disk" @@ -581,7 +587,8 @@ def show_conditions(self, cond, name=None, verbose=False): def print_source(cond, n): print(colored(n, 'cyan')) print(self.get_function_source(cond, n)) - + + _cond_path = settings.paths[cond] cond_names = list(self._functions[cond]) if name and name in cond_names and verbose: print_source(cond, name) @@ -606,9 +613,10 @@ def show_all_conds(self, cond, verbose=False): conditions. """ + _cond_path = settings.paths[cond] cond_names = [f[:-3] - for f in listdir(join(self._conds_path, cond)) - if isfile(join(join(self._conds_path, cond), f)) + for f in listdir(_cond_path) + if isfile(join(_cond_path, f)) and f != "__init__.py" and f[-3:] == ".py"] if verbose: diff --git a/polymorph/utils.py b/polymorph/utils.py index 87a771b..1191785 100644 --- a/polymorph/utils.py +++ b/polymorph/utils.py @@ -13,9 +13,42 @@ import platform import polymorph from os.path import dirname, join +import imp, os, importlib POLYM_PATH = dirname(polymorph.__file__) + +def import_file(filename, module=""): + """ + old stuff just for debug purposes + # m = importlib.import_module( + # "polymorph.conditions.%s.%s" % (settings.paths[cond], name)) + # importlib.reload(m) + + """ + module = None + + try: + # Get module name and path from full path + module_dir, module_file = os.path.split(filename) + module_name, module_ext = os.path.splitext(module_file) + + # Get module "spec" from filename + spec = importlib.util.spec_from_file_location(module_name,filename) + + module = spec.loader.load_module() + + except Exception as ec: + # TODO validate py2 importing if needed @shramos + print(ec) + + finally: + return module + + with open(filename, "r") as f: + return imp.load_module(module, f, filename, "") + + def capture(userfilter="", pcapname=".tmp.pcap", func=None, count=0, time=None): """This function is a wrapper function above the sniff scapy function. The result is a list of templates. The specification on filtering options can @@ -144,7 +177,7 @@ def set_ip_forwarding(value): file.write(str(value)) file.close() - + def get_arpspoofer(targets, gateway, iface=None, gatewaymac=None, ignore=None, arpmode='rep', mac=None, ip=None): # Creating a poison object @@ -154,4 +187,3 @@ def get_arpspoofer(targets, gateway, iface=None, gatewaymac=None, poisoner = ARPpoisoner(poison) # return the poisoner return poisoner -