Repository: google/docker-explorer

Issue:
Better detect overlay/overlay2 layouts
When processing an image that has both and "overlay2" and "overlay" folder it will take the "overlay2" folder by default even though there is no 'mounts' subfolder.

Discussion:

Patch:
diff --git a/docker_explorer/de.py b/docker_explorer/de.py
index ea8afa3..1e78818 100644
--- a/docker_explorer/de.py
+++ b/docker_explorer/de.py
@@ -24,7 +24,6 @@ import os
 
 from docker_explorer import errors
 from docker_explorer.lib import container
-from docker_explorer.lib import storage
 from docker_explorer.lib import utils
 
 
@@ -34,8 +33,6 @@ class DockerExplorer(object):
   Attributes:
     docker_directory (str): Path to use as the root of the Docker runtime.
       Default is '/var/lib/docker'.
-    storage_object (lib.Storage): The object implementing the methods for
-      exploring the Docker containers.
   """
 
   def __init__(self):
@@ -45,16 +42,12 @@ class DockerExplorer(object):
     self.containers_directory = None
     self.docker_directory = None
     self.docker_version = 2
-    self.storage_object = None
 
   def _SetDockerDirectory(self, docker_path):
     """Sets the Docker main directory.
 
     Args:
       docker_path(str): the absolute path to the docker directory.
-
-    Raises:
-      errors.BadStorageException: If the storage backend couldn't be detected.
     """
     self.docker_directory = docker_path
     if not os.path.isdir(self.docker_directory):
@@ -67,29 +60,6 @@ class DockerExplorer(object):
     self.containers_directory = os.path.join(
         self.docker_directory, 'containers')
 
-    if os.path.isfile(
-        os.path.join(self.docker_directory, 'repositories-aufs')):
-      # TODO: check this agains other storages in version 1.9 and below
-      self.docker_version = 1
-      self.storage_object = storage.AufsStorage(
-          docker_directory=self.docker_directory, docker_version=1)
-    elif os.path.isdir(os.path.join(self.docker_directory, 'overlay2')):
-      self.storage_object = storage.Overlay2Storage(
-          docker_directory=self.docker_directory)
-    elif os.path.isdir(os.path.join(self.docker_directory, 'overlay')):
-      self.storage_object = storage.OverlayStorage(
-          docker_directory=self.docker_directory)
-    elif os.path.isdir(os.path.join(self.docker_directory, 'aufs')):
-      self.storage_object = storage.AufsStorage(
-          docker_directory=self.docker_directory)
-    if self.storage_object is None:
-      err_message = (
-          'Could not detect storage system. '
-          'Make sure the docker directory ({0:s}) is correct. '
-          'If it is correct, you might want to run this script'
-          ' with higher privileges.'.format(self.docker_directory))
-      raise errors.BadStorageException(err_message)
-
   def AddBasicOptions(self, argument_parser):
     """Adds the global options to the argument_parser.
 
@@ -218,7 +188,7 @@ class DockerExplorer(object):
       mountpoint (str): the path to the destination mount point.
     """
     container_object = self.GetContainer(container_id)
-    self.storage_object.Mount(container_object, mountpoint)
+    container_object.Mount(mountpoint)
 
   def GetContainersString(self, only_running=False):
     """Returns a string describing the running containers.
@@ -260,10 +230,6 @@ class DockerExplorer(object):
     """
     print(self.GetContainersString(only_running=only_running))
 
-  def ShowRepositories(self):
-    """Displays information about the images in the Docker repository."""
-    print(self.storage_object.ShowRepositories())
-
   def ShowHistory(self, container_id, show_empty_layers=False):
     """Prints the modification history of a container.
 
@@ -274,6 +240,33 @@ class DockerExplorer(object):
     container_object = self.GetContainer(container_id)
     print(container_object.GetHistory(show_empty_layers))
 
+  def GetRepositoriesString(self):
+    """Returns information about images in the local Docker repositories.
+
+    Returns:
+      str: human readable list of images in local Docker repositories.
+    """
+    result_string = ''
+    repositories = []
+    if self.docker_version == 1:
+      repositories = [os.path.join(self.docker_directory, 'repositories-aufs')]
+    else:
+      image_path = os.path.join(self.docker_directory, 'image')
+      for storage_method in os.listdir(image_path):
+        repositories_file_path = os.path.join(
+            image_path, storage_method, 'repositories.json')
+        if os.path.isfile(repositories_file_path):
+          repositories.append(repositories_file_path)
+
+    for repositories_file_path in repositories:
+      result_string += (
+          'Listing repositories from file {0:s}\n'.format(
+              repositories_file_path))
+      with open(repositories_file_path) as rf:
+        result_string += utils.PrettyPrintJSON(rf.read())
+
+    return result_string
+
   def Main(self):
     """The main method for the DockerExplorer class.
 
@@ -287,7 +280,6 @@ class DockerExplorer(object):
 
     self._SetDockerDirectory(self.docker_directory)
 
-
     if options.command == 'mount':
       self.Mount(options.container_id, options.mountpoint)
 
@@ -301,7 +293,7 @@ class DockerExplorer(object):
       elif options.what == 'running_containers':
         self.ShowContainers(only_running=True)
       elif options.what == 'repositories':
-        self.ShowRepositories()
+        print(self.GetRepositoriesString())
 
     else:
       raise ValueError('Unhandled command %s' % options.command)
diff --git a/docker_explorer/lib/container.py b/docker_explorer/lib/container.py
index 4321385..04cec32 100644
--- a/docker_explorer/lib/container.py
+++ b/docker_explorer/lib/container.py
@@ -18,8 +18,10 @@ from __future__ import print_function, unicode_literals
 
 import json
 import os
+import subprocess
 
 from docker_explorer import errors
+from docker_explorer.lib import storage
 from docker_explorer.lib import utils
 
 
@@ -38,11 +40,18 @@ class Container(object):
     name (str): the name of the container.
     running (boolean): True if the container is running.
     start_timestamp (str): the container's start timestamp.
-    storage_driver (str): the container's storage driver.
+    storage_name (str): the container's storage driver name.
+    storage_object (BaseStorage): the container's storage backend object.
     volumes (list(tuple)): list of mount points to bind from host to the
       container. (Docker storage backend v1).
   """
 
+  STORAGES_MAP = {
+      'aufs': storage.AufsStorage,
+      'overlay': storage.OverlayStorage,
+      'overlay2': storage.Overlay2Storage
+  }
+
   def __init__(self, docker_directory, container_id, docker_version=2):
     """Initializes the Container class.
 
@@ -88,16 +97,18 @@ class Container(object):
     if json_state:
       self.running = json_state.get('Running', False)
       self.start_timestamp = json_state.get('StartedAt', False)
-    self.storage_driver = container_info_dict.get('Driver', None)
-    if self.storage_driver is None:
+    self.storage_name = container_info_dict.get('Driver', None)
+    if self.storage_name is None:
       raise errors.BadContainerException(
           '{0} container config file lacks Driver key'.format(
               container_info_json_path))
+
+    self._SetStorage(self.storage_name)
     self.volumes = container_info_dict.get('Volumes', None)
 
     if self.docker_version == 2:
       c_path = os.path.join(
-          self.docker_directory, 'image', self.storage_driver, 'layerdb',
+          self.docker_directory, 'image', self.storage_name, 'layerdb',
           'mounts', container_id)
       with open(os.path.join(c_path, 'mount-id')) as mount_id_file:
         self.mount_id = mount_id_file.read()
@@ -131,7 +142,7 @@ class Container(object):
     elif self.docker_version == 2:
       hash_method, layer_id = layer_id.split(':')
       layer_info_path = os.path.join(
-          self.docker_directory, 'image', self.storage_driver, 'imagedb',
+          self.docker_directory, 'image', self.storage_name, 'imagedb',
           'content', hash_method, layer_id)
     if os.path.isfile(layer_info_path):
       with open(layer_info_path) as layer_info_file:
@@ -162,7 +173,7 @@ class Container(object):
       elif self.docker_version == 2:
         hash_method, layer_id = current_layer.split(':')
         parent_layer_path = os.path.join(
-            self.docker_directory, 'image', self.storage_driver, 'imagedb',
+            self.docker_directory, 'image', self.storage_name, 'imagedb',
             'metadata', hash_method, layer_id, 'parent')
         if not os.path.isfile(parent_layer_path):
           break
@@ -204,3 +215,41 @@ class Container(object):
         else:
           history_str += 'Empty layer'
     return history_str
+
+  def _SetStorage(self, storage_name):
+    """Sets the storage_object attribute.
+
+    Args:
+      storage_name (str): the name of the storage.
+    Returns:
+      BaseStorage: a storage object.
+    Raises:
+      BadContainerException: if no storage Driver is defined, or if it is not
+        implemented
+    """
+    storage_class = self.STORAGES_MAP.get(storage_name, None)
+
+    if storage_class is None:
+      raise errors.BadContainerException(
+          'Storage driver {0} is not implemented'.format(storage_name))
+
+    self.storage_object = storage_class(
+        self.docker_directory, self.docker_version)
+
+  def Mount(self, mount_dir):
+    """Mounts the specified container's filesystem.
+
+    Args:
+      mount_dir (str): the path to the destination mount point
+    """
+
+    commands = self.storage_object.MakeMountCommands(self, mount_dir)
+    for c in commands:
+      print(c)
+    print('Do you want to mount this container ID: {0:s} on {1:s} ?\n'
+          '(ie: run these commands) [Y/n]'.format(self.container_id, mount_dir))
+    choice = raw_input().lower()
+    if not choice in ['y', 'yes', '']:
+      for c in commands:
+        # TODO() this is quite unsafe, need to properly split args
+        subprocess.call(c, shell=True)
diff --git a/docker_explorer/lib/storage.py b/docker_explorer/lib/storage.py
index 1735fa2..2cc96c1 100644
--- a/docker_explorer/lib/storage.py
+++ b/docker_explorer/lib/storage.py
@@ -17,11 +17,8 @@
 from __future__ import print_function, unicode_literals
 
 import os
-import subprocess
 import sys
 
-from docker_explorer.lib import utils
-
 
 class BaseStorage(object):
   """This class provides tools to list and access containers metadata.
@@ -51,24 +48,6 @@ class BaseStorage(object):
     if self.docker_version == 1:
       self.container_config_filename = 'config.json'
 
-  def ShowRepositories(self):
-    """Returns information about the images in the Docker repository.
-
-    Returns:
-      str: human readable information about image repositories.
-    """
-    repositories_file_path = os.path.join(
-        self.docker_directory, 'image', self.STORAGE_METHOD,
-        'repositories.json')
-    if self.docker_version == 1:
-      repositories_file_path = os.path.join(
-          self.docker_directory, 'repositories-aufs')
-    result_string = (
-        'Listing repositories from file {0:s}').format(repositories_file_path)
-    with open(repositories_file_path) as rf:
-      repositories_string = rf.read()
-    return result_string + utils.PrettyPrintJSON(repositories_string)
-
   def MakeMountCommands(self, container_object, mount_dir):
     """Generates the required shell commands to mount a container given its ID.
 
@@ -123,25 +102,6 @@ class BaseStorage(object):
 
     return extra_commands
 
-  def Mount(self, container_object, mount_dir):
-    """Mounts the specified container's filesystem.
-
-    Args:
-      container_object (Container): the container.
-      mount_dir (str): the path to the destination mount point
-    """
-
-    commands = self.MakeMountCommands(container_object, mount_dir)
-    for c in commands:
-      print(c)
-    print('Do you want to mount this container Id: {0:s} on {1:s} ?\n'
-          '(ie: run these commands) [Y/n]'.format(
-              container_object.container_id, mount_dir))
-    choice = raw_input().lower()
-    if not choice or choice == 'y' or choice == 'yes':
-      for c in commands:
-        # TODO(romaing) this is quite unsafe, need to properly split args
-        subprocess.call(c, shell=True)
 
 class AufsStorage(BaseStorage):
   """This class implements AuFS storage specific methods."""
diff --git a/docker_explorer/lib/utils.py b/docker_explorer/lib/utils.py
index b4e9db3..cdfb2b8 100644
--- a/docker_explorer/lib/utils.py
+++ b/docker_explorer/lib/utils.py
@@ -44,5 +44,6 @@ def PrettyPrintJSON(string):
   Returns:
     str: pretty printed JSON string.
   """
-  return json.dumps(
+  pretty_json = json.dumps(
       json.loads(string), sort_keys=True, indent=4, separators=(', ', ': '))
+  return pretty_json + '\n'


Test patch:
diff --git a/tests.py b/tests.py
index f7a79f4..99ca551 100644
--- a/tests.py
+++ b/tests.py
@@ -47,7 +47,7 @@ class UtilsTests(unittest.TestCase):
     test_json = json.dumps(test_dict)
     expected_string = ('{\n    "test": [\n        {\n            "dict1": {\n'
                        '                "key1": "val1"\n            }, \n'
-                       '            "dict2": null\n        }\n    ]\n}')
+                       '            "dict2": null\n        }\n    ]\n}\n')
     self.assertEqual(expected_string, utils.PrettyPrintJSON(test_json))
 
 
@@ -89,8 +89,8 @@ class TestDEMain(unittest.TestCase):
     self.assertEqual(expected_error_message, err.exception.message)
 
 
-class StorageTestCase(unittest.TestCase):
-  """Base class for tests of different BaseStorage implementations."""
+class DockerTestCase(unittest.TestCase):
+  """Base class for tests of different Storage implementations."""
 
   @classmethod
   def tearDownClass(cls):
@@ -98,7 +98,7 @@ class StorageTestCase(unittest.TestCase):
 
   @classmethod
   def _setup(cls, driver, driver_class):
-    """Internal method to set up the TestCase on a specific storate."""
+    """Internal method to set up the TestCase on a specific storage."""
     cls.driver = driver
     docker_directory_path = os.path.join('test_data', 'docker')
     if not os.path.isdir(docker_directory_path):
@@ -113,17 +113,17 @@ class StorageTestCase(unittest.TestCase):
 
   def testDetectStorage(self):
     """Tests the DockerExplorer.DetectStorage function."""
-    storage_object = self.de_object.storage_object
-    self.assertIsNotNone(storage_object)
-    self.assertIsInstance(storage_object, self.driver_class)
-    self.assertEqual(storage_object.STORAGE_METHOD, self.driver)
+    for container_obj in self.de_object.GetAllContainers():
+      self.assertIsNotNone(container_obj.storage_object)
+      self.assertEqual(container_obj.storage_name, self.driver)
+      self.assertIsInstance(container_obj.storage_object, self.driver_class)
 
-    self.assertEqual(2, storage_object.docker_version)
-    self.assertEqual('config.v2.json',
-                     self.de_object.container_config_filename)
+      self.assertEqual(2, container_obj.docker_version)
+      self.assertEqual(
+          'config.v2.json', container_obj.container_config_filename)
 
 
-class TestAufsStorage(StorageTestCase):
+class TestAufsStorage(DockerTestCase):
   """Tests methods in the BaseStorage object."""
 
   @classmethod
@@ -198,12 +198,13 @@ class TestAufsStorage(StorageTestCase):
     self.assertEqual(['/bin/sh', '-c', '#(nop) ', 'CMD ["sh"]'],
                      layer_info['container_config']['Cmd'])
 
-  def testShowRepositories(self):
-    """Tests the BaseStorage.ShowRepositories function on a AUFS storage."""
-    result_string = self.de_object.storage_object.ShowRepositories()
+  def testGetRepositoriesString(self):
+    """Tests BaseStorage.GetRepositoriesString() on a AUFS storage."""
+    self.maxDiff = None
+    result_string = self.de_object.GetRepositoriesString()
     expected_string = (
         'Listing repositories from file '
-        'test_data/docker/image/aufs/repositories.json{\n'
+        'test_data/docker/image/aufs/repositories.json\n{\n'
         '    "Repositories": {\n'
         '        "busybox": {\n'
         '            "busybox:latest": '
@@ -211,7 +212,7 @@ class TestAufsStorage(StorageTestCase):
         '68"\n'
         '        }\n'
         '    }\n'
-        '}')
+        '}\n')
     self.assertEqual(expected_string, result_string)
 
   def testMakeMountCommands(self):
@@ -219,7 +220,7 @@ class TestAufsStorage(StorageTestCase):
     container_id = (
         '7b02fb3e8a665a63e32b909af5babb7d6ba0b64e10003b2d9534c7d5f2af8966')
     container_obj = self.de_object.GetContainer(container_id)
-    commands = self.de_object.storage_object.MakeMountCommands(
+    commands = container_obj.storage_object.MakeMountCommands(
         container_obj, '/mnt')
     expected_commands = [
         ('mount -t aufs -o ro,br=test_data/docker/aufs/diff/test_data/docker/'
@@ -253,7 +254,7 @@ class TestAufsStorage(StorageTestCase):
     self.assertEqual(expected_string, container_obj.GetHistory())
 
 
-class TestOverlayStorage(StorageTestCase):
+class TestOverlayStorage(DockerTestCase):
   """Tests methods in the OverlayStorage object."""
 
   @classmethod
@@ -329,13 +330,13 @@ class TestOverlayStorage(StorageTestCase):
     self.assertEqual(['/bin/sh', '-c', '#(nop) ', 'CMD ["sh"]'],
                      layer_info['container_config']['Cmd'])
 
-  def testShowRepositories(self):
-    """Tests the BaseStorage.ShowRepositories function on a Overlay storage."""
-    result_string = self.de_object.storage_object.ShowRepositories()
+  def testGetRepositoriesString(self):
+    """Tests BaseStorage.GetRepositoriesString() on a Overlay storage."""
+    result_string = self.de_object.GetRepositoriesString()
     self.maxDiff = None
     expected_string = (
         'Listing repositories from file '
-        'test_data/docker/image/overlay/repositories.json{\n'
+        'test_data/docker/image/overlay/repositories.json\n{\n'
         '    "Repositories": {\n'
         '        "busybox": {\n'
         '            "busybox:latest": "sha256:'
@@ -346,7 +347,7 @@ class TestOverlayStorage(StorageTestCase):
         '2c3"\n'
         '        }\n'
         '    }\n'
-        '}')
+        '}\n')
     self.assertEqual(expected_string, result_string)
 
   def testMakeMountCommands(self):
@@ -354,7 +355,7 @@ class TestOverlayStorage(StorageTestCase):
     container_id = (
         '5dc287aa80b460652a5584e80a5c8c1233b0c0691972d75424cf5250b917600a')
     container_obj = self.de_object.GetContainer(container_id)
-    commands = self.de_object.storage_object.MakeMountCommands(
+    commands = container_obj.storage_object.MakeMountCommands(
         container_obj, '/mnt')
     expected_commands = [(
         'mount -t overlay overlay -o ro,lowerdir='
@@ -381,7 +382,7 @@ class TestOverlayStorage(StorageTestCase):
     self.assertEqual(expected_string, container_obj.GetHistory())
 
 
-class TestOverlay2Storage(StorageTestCase):
+class TestOverlay2Storage(DockerTestCase):
   """Tests methods in the Overlay2Storage object."""
 
   @classmethod
@@ -457,13 +458,13 @@ class TestOverlay2Storage(StorageTestCase):
     self.assertEqual(['/bin/sh', '-c', '#(nop) ', 'CMD ["sh"]'],
                      layer_info['container_config']['Cmd'])
 
-  def testShowRepositories(self):
-    """Tests the BaseStorage.ShowRepositories function on a Overlay2 storage."""
-    result_string = self.de_object.storage_object.ShowRepositories()
+  def testGetRepositoriesString(self):
+    """Tests BaseStorage.GetRepositoriesString() on a Overlay2 storage."""
+    result_string = self.de_object.GetRepositoriesString()
     self.maxDiff = None
     expected_string = (
         'Listing repositories from file '
-        'test_data/docker/image/overlay2/repositories.json{\n'
+        'test_data/docker/image/overlay2/repositories.json\n{\n'
         '    "Repositories": {\n'
         '        "busybox": {\n'
         '            "busybox:latest": "sha256:'
@@ -474,7 +475,14 @@ class TestOverlay2Storage(StorageTestCase):
         'c7"\n'
         '        }\n'
         '    }\n'
-        '}')
+        '}\n'
+        'Listing repositories from file '
+        'test_data/docker/image/overlay/repositories.json\n'
+        '{\n'
+        '    "Repositories": {}\n'
+        '}\n'
+
+    )
     self.assertEqual(expected_string, result_string)
 
   def testMakeMountCommands(self):
@@ -483,7 +491,7 @@ class TestOverlay2Storage(StorageTestCase):
     container_id = (
         '8e8b7f23eb7cbd4dfe7e91646ddd0e0f524218e25d50113559f078dfb2690206')
     container_obj = self.de_object.GetContainer(container_id)
-    commands = self.de_object.storage_object.MakeMountCommands(
+    commands = container_obj.storage_object.MakeMountCommands(
         container_obj, '/mnt')
     expected_commands = [(
         'mount -t overlay overlay -o ro,lowerdir='
@@ -511,7 +519,7 @@ class TestOverlay2Storage(StorageTestCase):
         'with command : /bin/sh -c #(nop)  CMD ["sh"]')
     self.assertEqual(expected_string, container_obj.GetHistory(container_obj))
 
-del StorageTestCase
+del DockerTestCase
 
 if __name__ == '__main__':
   unittest.main()