large refactor; Moved module level code into class; added docstrings, exception handling, and more tests;

Author: cnelson

.gitignore

@@ -3,3 +3,4 @@ dist
 *.egg*
 *.pyc
 .tox
+.coverage

nsenter/__init__.py

@@ -6,6 +6,7 @@
 
 import argparse
 import ctypes
+import errno
 import os
 import logging
 from pathlib import Path
@@ -14,66 +15,133 @@
 except ImportError:
     from contextlib2 import ExitStack
 
-NAMESPACE_NAMES = frozenset('mnt ipc net pid user uts'.split())
+NAMESPACE_NAMES = frozenset(['mnt','ipc','net','pid','user','uts'])
 
-log = logging.getLogger('nsenter')
 
-libc = ctypes.CDLL('libc.so.6')
+class Namespace(object):
+    """A context manager for entering namespaces
 
+    Args:
+        pid: The PID for the owner of the namespace to enter
+        ns_type: The type of namespace to enter must be one of
+            mnt ipc net pid user uts
 
-def nsfd(process, ns_type):
-    """
-    Returns the namespace file descriptor for process (self or PID) and namespace type
-    """
-    return Path('/proc') / str(process) / 'ns' / ns_type
+    Raises:
+        IOError: A non existent PID was provided
+        ValueError: An improper ns_type was provided
+        OSError: Unable to enter or exit the namespace
 
-nsfd.__annotations__ = {'process': str, 'ns_type': str, 'return': Path}
+    Example:
+        with Namespace(<pid>, <ns_type>):
+            #do something in the namespace
+            pass
+    """
 
+    _log = logging.getLogger(__name__)
+    _libc = ctypes.CDLL('libc.so.6', use_errno=True)
 
-class Namespace(object):
     def __init__(self, pid, ns_type):
+        if ns_type not in NAMESPACE_NAMES:
+            raise ValueError('ns_type must be one of {0}'.format(
+                ', '.join(NAMESPACE_NAMES)
+            ))
+
         self.pid = pid
         self.ns_type = ns_type
-        self.parent_fd = nsfd('self', ns_type).open()
-        self.parent_fileno = self.parent_fd.fileno()
-        self.target_fd = nsfd(pid, ns_type).open()
+
+        self.target_fd = self._nsfd(pid, ns_type).open()
         self.target_fileno = self.target_fd.fileno()
 
-    __init__.__annotations__ = {'pid': str, 'ns_type': str}
+        self.parent_fd = self._nsfd('self', ns_type).open()
+        self.parent_fileno = self.parent_fd.fileno()
 
-    def __enter__(self):
-        log.debug('Entering %s namespace %s', self.ns_type, self.pid)
-        libc.setns(self.target_fileno, 0)
+    __init__.__annotations__  = {'pid': str, 'ns_type': str}
 
-    def __exit__(self, type, value, tb):
-        log.debug('Leaving %s namespace %s', self.ns_type, self.pid)
-        libc.setns(self.parent_fileno, 0)
+    def _nsfd(self, pid, ns_type):
+        """Utility method to build a pathlib.Path instance pointing at the 
+        requested namespace entry
+
+        Args:
+            pid: The PID 
+            ns_type: The namespace type to enter
+
+        Returns:
+             pathlib.Path pointing to the /proc namespace entry
+        """
+        return Path('/proc') / str(pid) / 'ns' / ns_type
+
+    _nsfd.__annotations__ = {'process': str, 'ns_type': str, 'return': Path}
+
+    def _close_files(self):
+        """Utility method to close our open file handles"""
         try:
             self.target_fd.close()
         except:
             pass
-        self.parent_fd.close()
+        self.parent_fd.close()    
+
+    def __enter__(self):
+        self._log.debug('Entering %s namespace %s', self.ns_type, self.pid)
 
+        if self._libc.setns(self.target_fileno, 0) == -1:
+            e = ctypes.get_errno()
+            self._close_files()
+            raise OSError(e, errno.errorcode[e]) 
 
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
+    def __exit__(self, type, value, tb):
+        self._log.debug('Leaving %s namespace %s', self.ns_type, self.pid)
+        
+        if self._libc.setns(self.parent_fileno, 0) == -1:
+            e = ctypes.get_errno()
+            self._close_files()
+            raise OSError(e, errno.errorcode[e]) 
+
+        self._close_files()
+        
+def main(): #pragma: no cover
+    """Command line interface to the Namespace Contet Manager"""
+
+    parser = argparse.ArgumentParser(prog='nsenter', description=__doc__)
+    
     parser.add_argument('--target', '-t', required=True, metavar='PID',
-                        help='Specify a target process to get contexts from.')
+                        help='A target process to get contexts from')
+
+    group = parser.add_argument_group('Namespaces')
+
     for ns in NAMESPACE_NAMES:
-        parser.add_argument('--{0}'.format(ns), action='store_true', help='Enter the {0} namespace'.format(ns))
-    parser.add_argument('--all', action='store_true', help='Enter all namespaces')
+       group.add_argument('--{0}'.format(ns), action='store_true', 
+                            help='Enter the {0} namespace'.format(ns))
+    
+    parser.add_argument('--all', action='store_true', 
+                        help='Enter all namespaces')
+
     parser.add_argument('command', nargs='*', default=['/bin/sh'])
 
     args = parser.parse_args()
 
-    with ExitStack() as stack:
-        namespaces = []
-        for ns in NAMESPACE_NAMES:
-            if getattr(args, ns) or args.all:
-                namespaces.append(Namespace(args.target, ns))
-        for ns in namespaces:
-            stack.enter_context(ns)
-        os.execl(args.command[0], *args.command)
-
-if __name__ == '__main__':
+    #make sure we have --all or at least one namespace
+    if (True not in [getattr(args, ns) for ns in NAMESPACE_NAMES]
+        and not args.all):
+        parser.error('You must specify at least one namespace')
+
+    try:
+        with ExitStack() as stack:
+            namespaces = []
+            for ns in NAMESPACE_NAMES:
+                if getattr(args, ns) or args.all:
+                    namespaces.append(Namespace(args.target, ns))
+           
+            for ns in namespaces:
+                 stack.enter_context(ns)
+
+            os.execlp(args.command[0], *args.command)
+    except IOError as exc:
+        parser.error('Unable to access PID: {0}'.format(exc))
+    except OSError as exc:
+        parser.error('Unable to enter {0} namespace: {1}'.format(
+            ns.ns_type, exc
+        ))
+
+
+if __name__ == '__main__': #pragma: no cover
     main()

tests.py

@@ -1,24 +1,70 @@
 import unittest
 
-import subprocess
+import subprocess, os, errno
 
 from nsenter import Namespace, NAMESPACE_NAMES 
 
 class TestNamespaces(unittest.TestCase):
 
     def setUp(self):
+        """Spawn a child process so we have a PID to enter"""
+        
         self._child = subprocess.Popen(['/bin/cat'])
 
     def tearDown(self):
+        """SIGTERM the child process"""
+        
         self._child.terminate()
         self._child.wait()
 
-    def test_namespace(self):
-        for name in NAMESPACE_NAMES: 
-            with Namespace(self._child.pid, name):
-                assert True
+    def test_namespaces_except_user(self):
+        """Test entering all namespaces execept user
+
+        Must have CAP_SYS_ADMIN to run these tests properly
+        """
+
+        #Can't use the assertRaises context manager in python2.6
+        def do_test():
+            for name in filter(lambda x: x != 'user', NAMESPACE_NAMES):
+                with Namespace(self._child.pid, name):
+                    pass
+            
+            #if we aren't root (technically: CAP_SYS_ADMIN)
+            #then we'll get OSError (EPERM) for all our tests
+            if os.geteuid() != 0:
+               self.assertRaises(OSError, do_test)
+            else:
+                do_test()
+
+    def test_user_namespace(self):
+        """Test entering a non-existent namespace"""
+        
+        def do_test():
+            with Namespace(self._child.pid, 'user'):
+                pass
+
+        #This process doesn't have a user namespace
+        #So this will OSError(EINVAL)
+        self.assertRaises(OSError, do_test)
+
+    def test_bad_namespace(self):
+        """Test entering a bad namespace type"""
+        
+        def do_test():
+            with Namespace(self._child.pid, 'foo'):
+                pass
+        self.assertRaises(ValueError, do_test)
+
+    def test_bad_pid(self):
+        """Test entering bad pid's name space"""
+        
+        def do_test():
+            with Namespace('foo', 'net'):
+                pass
+
+        self.assertRaises(IOError, do_test)
+
 
-        pass
 
 if __name__ == '__main__':
     unittest.main()