Add experimental monkey patching

2 files changed, 39 insertions, 21 deletions

diff --git a/defusedxml/__init__.py b/defusedxml/__init__.py
index 1cc84c9..4c3b996 100644
--- a/defusedxml/__init__.py
+++ b/defusedxml/__init__.py
@@ -8,16 +8,35 @@
 from __future__ import print_function, absolute_import
 
 from .common import (DefusedXmlException, DTDForbidden, EntitiesForbidden,
-                     ExternalReferenceForbidden, NotSupportedError)
-
-#from . import cElementTree
-#from . import ElementTree
-#from . import minidom
-#from . import pulldom
-#from . import sax
-#from . import xmlrpc
-#from . import expatbuilder
-#from . import expatreader
+                     ExternalReferenceForbidden, NotSupportedError,
+                     _apply_defusing)
+
+def defuse_stdlib():
+    """Monkey patch and defuse all stdlib packages
+
+    :warning: It's an HIGHLY EXPERIMETNAL and hardly tested feature.
+    """
+    defused = {}
+
+    from . import cElementTree
+    from . import ElementTree
+    from . import minidom
+    from . import pulldom
+    from . import sax
+    from . import expatbuilder
+    from . import expatreader
+    from . import xmlrpc
+
+    xmlrpc.monkey_patch()
+    defused[xmlrpc] = None
+
+    for defused_mod in [cElementTree, ElementTree, minidom, pulldom, sax,
+                        expatbuilder, expatreader]:
+        stdlib_mod = _apply_defusing(defused_mod)
+        defused[defused_mod] = stdlib_mod
+
+    return defused
+
 
 __version__ = "0.3dev"
 
diff --git a/defusedxml/common.py b/defusedxml/common.py
index cf1edff..5e5f8a2 100644
--- a/defusedxml/common.py
+++ b/defusedxml/common.py
@@ -71,18 +71,17 @@ class NotSupportedError(DefusedXmlException):
     """
 
 
-def _wire_module(srcmod, dstmodname):
-    assert srcmod is sys.modules[srcmod.__name__]
-    dstmod = sys.modules[dstmodname]
-    names = getattr(srcmod, "__all__", None)
-    if not names:
-        names = tuple(name for name in dir(srcmod)
-                      if not name.startswith("_"))
-    for name in names:
-        if hasattr(dstmod, name):
+def _apply_defusing(defused_mod):
+    assert defused_mod is sys.modules[defused_mod.__name__]
+    stdlib_name = defused_mod.__origin__
+    __import__(stdlib_name, {}, {}, ["*"])
+    stdlib_mod = sys.modules[stdlib_name]
+    stdlib_names = set(dir(stdlib_mod))
+    for name, obj in vars(defused_mod).items():
+        if name.startswith("_") or name not in stdlib_names:
             continue
-        value = getattr(srcmod, name)
-        setattr(dstmod, name, value)
+        setattr(stdlib_mod, name, obj)
+    return stdlib_mod
 
 
 def _generate_etree_functions(DefusedXMLParser, _TreeBuilder,