diff --git a/envs.py b/envs.py index 0fd0db4..03aa67a 100644 --- a/envs.py +++ b/envs.py @@ -1,6 +1,5 @@ # -*- coding: utf-8 -*- - from __future__ import absolute_import from __future__ import print_function @@ -8,39 +7,44 @@ import sys import idaapi - -def root_path(): - return os.path.abspath(os.sep) - - def detect_env(): + """ + Detect and load a Python virtual environment if present. + """ + print("Detecting python virtual environment...") + try: - print("Detecting Virtualenv envs") path = activate_virtualenv_env(None, False) - print("Virtualenv env activated: " + str(path)) + print(" - Virtualenv activated: %s" % path) return True except ValueError as e: - print(e.message) + print(" - Virtualenv: %s" % e) + try: - print("Detecting Conda envs") path = activate_conda_env(None, None, False) - print("Conda env activated: " + str(path)) + print(" - Conda env activated: %s" % path) return True except ValueError as e: - print(e.message) - return False + print(" - Conda: %s" % e) + return False def activate_virtualenv_env(virtualenv=None, interactive=True): + """ + Activate a Virtualenv based virtual environment. + """ folder = "Scripts" if os.name == "nt" else "bin" + if virtualenv == None: virtualenv = os.environ.get("VIRTUAL_ENV") + if not virtualenv and interactive: default_virtualenv = os.path.join(idaapi.get_user_idadir(), "virtualenv") virtualenv = idaapi.askstr(0, default_virtualenv, "Select a virtualenv") if not virtualenv: - raise ValueError("No virtualenv env") + raise ValueError("No active virtualenv") + if not os.path.isdir(virtualenv): raise ValueError("This path is not a dir: " + virtualenv) @@ -48,23 +52,27 @@ def activate_virtualenv_env(virtualenv=None, interactive=True): if not os.path.isfile(virtualenv_script): raise ValueError('Enable to find "' + folder + os.sep + 'activate_this.py" in virtualenv: ' + virtualenv) - execfile(virtualenv_script, dict(__file__=virtualenv_script)) + with open(virtualenv_script) as infile: + exec(infile.read(), dict(__file__=virtualenv_script)) return virtualenv - def activate_conda_env(base=None, env=None, interactive=True): + """ + Activate a Conda based virtual environment. + """ folder = "Scripts" if os.name == "nt" else "bin" # Get env if env == None: env = os.environ.get("CONDA_PREFIX") if not env and interactive: - env = idaapi.askstr(0, root_path(), "Select a env") + root_path = os.path.abspath(os.sep) + env = idaapi.askstr(0, root_path, "Select a env") # Check env if not env: - raise ValueError("No Conda env") + raise ValueError("No active Conda env") if not os.path.isdir(env): raise ValueError("This path is not a dir: " + env) @@ -88,6 +96,7 @@ def activate_conda_env(base=None, env=None, interactive=True): site.addsitedir(site_packages) sys.real_prefix = sys.prefix sys.prefix = env + # Move the added items to the front of the path: new_sys_path = [] for item in list(sys.path): @@ -97,3 +106,7 @@ def activate_conda_env(base=None, env=None, interactive=True): sys.path[:0] = new_sys_path return env + +# make the env activation functions accessible to the IDA console (7.0+) +sys.modules["__main__"].activate_virtualenv_env = activate_virtualenv_env +sys.modules["__main__"].activate_conda_env = activate_conda_env \ No newline at end of file