From 824361c7148b43dc3025831170094ff17ba2a7bd Mon Sep 17 00:00:00 2001
From: Arcadiy Ivanov <arcadiy@ivanov.biz>
Date: Tue, 3 Apr 2018 19:07:38 -0400
Subject: [PATCH] Docker rebase

Allows to rebase docker image on another parent. v2 support only (no means to test on v1).
---
 .travis.yml                 |  3 ++
 Makefile                    | 11 ++++++
 circle.yml                  |  9 ++++-
 docker_squash/cli.py        | 36 ++++++++++-------
 docker_squash/image.py      | 51 ++++++++++++++++--------
 docker_squash/squash.py     |  7 ++--
 docker_squash/v1_image.py   |  2 +-
 docker_squash/v2_image.py   | 79 +++++++++++++++++++++++++++----------
 tests/test_integ_squash.py  | 75 +++++++++++++++++++++++++----------
 tests/test_unit_squash.py   |  2 +-
 tests/test_unit_v2_image.py | 28 ++++++++++++-
 tox.ini                     |  2 +-
 12 files changed, 223 insertions(+), 82 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 6cfef6e..5f9c68d 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -2,7 +2,10 @@ language: python
 python:
   - "2.7"
   - "3.4"
+  - "3.5"
+  - "3.6"
   - "pypy"
+  - "pypy3"
 install: "pip install -r requirements.txt"
 script:
   - "py.test -v"
diff --git a/Makefile b/Makefile
index f3261ba..ae064eb 100644
--- a/Makefile
+++ b/Makefile
@@ -12,6 +12,12 @@ test-py34: prepare
 test-py35: prepare
 	tox -e py35 -- tests
 
+test-py36: prepare
+	tox -e py36 -- tests
+
+test-py37: prepare
+	tox -e py37 -- tests
+
 test-unit: prepare
 	tox -- tests/test_unit*
 
@@ -28,6 +34,11 @@ else
 	@sudo chmod +x /usr/bin/docker
 endif
 
+ci-install-pythons:
+	curl -L https://github.com/pyenv/pyenv-installer/raw/master/bin/pyenv-installer | bash
+	pyenv update
+	for pyver in 2.7.14 3.4.8 3.5.5 3.6.5; do pyenv install -s $$pyver; done
+
 ci-publish-junit:
 	@mkdir -p ${CIRCLE_TEST_REPORTS}
 	@cp target/junit*.xml ${CIRCLE_TEST_REPORTS}
diff --git a/circle.yml b/circle.yml
index aecf08c..1023e3b 100644
--- a/circle.yml
+++ b/circle.yml
@@ -12,9 +12,14 @@ machine:
     CI: true
 
 dependencies:
+  pre:
+    - mkdir -p ~/.pyenv
+    - make -f Makefile ci-install-pythons
+  cache_directories:
+    - ~/.pyenv
   override:
     - pip install tox tox-pyenv docker-py>=1.7.2 six
-    - pyenv local 2.7.11 3.4.4 3.5.1
+    - pyenv local 2.7.14 3.4.8 3.5.5 3.6.5
   post:
     - docker version
     - docker info
@@ -27,5 +32,7 @@ test:
         parallel: true
     - make test-py35:
         parallel: true
+    - make test-py36:
+        parallel: true
   post:
     - make ci-publish-junit
diff --git a/docker_squash/cli.py b/docker_squash/cli.py
index 267df25..fb48ce9 100644
--- a/docker_squash/cli.py
+++ b/docker_squash/cli.py
@@ -60,18 +60,22 @@ def run(self):
             '--version', action='version', help='Show version and exit', version=version)
 
         parser.add_argument('image', help='Image to be squashed')
-        parser.add_argument(
-            '-d', '--development', action='store_true', help='Does not clean up after failure for easier debugging')
-        parser.add_argument(
-            '-f', '--from-layer', help='ID of the layer or image ID or image name. If not specified will squash all layers in the image')
-        parser.add_argument(
-            '-t', '--tag', help="Specify the tag to be used for the new image. If not specified no tag will be applied")
-        parser.add_argument(
-            '-c', '--cleanup', action='store_true', help="Remove source image from Docker after squashing")
-        parser.add_argument(
-            '--tmp-dir', help='Temporary directory to be created and used')
-        parser.add_argument(
-            '--output-path', help='Path where the image should be stored after squashing. If not provided, image will be loaded into Docker daemon')
+        parser.add_argument('-r', '--rebase',
+                            help='Rebase the image on a different "FROM"')
+        parser.add_argument('-d', '--development', action='store_true',
+                            help='Does not clean up after failure for easier debugging')
+        parser.add_argument('-f', '--from-layer',
+                            help='ID of the layer or image ID or image name. '
+                                 'If not specified will squash all layers in the image')
+        parser.add_argument('-t', '--tag',
+                            help="Specify the tag to be used for the new image. If not specified no tag will be applied")
+        parser.add_argument('-c', '--cleanup', action='store_true',
+                            help="Remove source image from Docker after squashing")
+        parser.add_argument('--tmp-dir',
+                            help='Temporary directory to be created and used')
+        parser.add_argument('--output-path',
+                            help='Path where the image should be stored after squashing. '
+                                 'If not provided, image will be loaded into Docker daemon')
 
         args = parser.parse_args()
 
@@ -84,7 +88,8 @@ def run(self):
 
         try:
             squash.Squash(log=self.log, image=args.image,
-                          from_layer=args.from_layer, tag=args.tag, output_path=args.output_path, tmp_dir=args.tmp_dir, development=args.development, cleanup=args.cleanup).run()
+                          from_layer=args.from_layer, tag=args.tag, output_path=args.output_path, tmp_dir=args.tmp_dir,
+                          development=args.development, cleanup=args.cleanup, rebase=args.rebase).run()
         except KeyboardInterrupt:
             self.log.error("Program interrupted by user, exiting...")
             sys.exit(1)
@@ -96,8 +101,9 @@ def run(self):
             else:
                 self.log.error(str(e))
 
-            self.log.error(
-                "Execution failed, consult logs above. If you think this is our fault, please file an issue: https://github.com/goldmann/docker-squash/issues, thanks!")
+            self.log.error("Execution failed, consult logs above. "
+                           "If you think this is our fault, please file an issue: "
+                           "https://github.com/goldmann/docker-squash/issues, thanks!")
 
             if isinstance(e, SquashError):
                 sys.exit(e.code)
diff --git a/docker_squash/image.py b/docker_squash/image.py
index 5616026..773d5c8 100644
--- a/docker_squash/image.py
+++ b/docker_squash/image.py
@@ -42,17 +42,19 @@ class Image(object):
     FORMAT = None
     """ Image format version """
 
-    def __init__(self, log, docker, image, from_layer, tmp_dir=None, tag=None):
+    def __init__(self, log, docker, image, from_layer, tmp_dir=None, tag=None, rebase=None):
         self.log = log
         self.debug = self.log.isEnabledFor(logging.DEBUG)
         self.docker = docker
         self.image = image
         self.from_layer = from_layer
         self.tag = tag
+        self.rebase = rebase
         self.image_name = None
         self.image_tag = None
         self.squash_id = None
 
+
         # Workaround for https://play.golang.org/p/sCsWMXYxqy
         #
         # Golang doesn't add padding to microseconds when marshaling
@@ -69,6 +71,7 @@ def __init__(self, log, docker, image, from_layer, tmp_dir=None, tag=None):
 
     def squash(self):
         self._before_squashing()
+        self.log.info("Squashing image '%s'..." % self.image)
         ret = self._squash()
         self._after_squashing()
 
@@ -92,12 +95,14 @@ def _initialize_directories(self):
 
         # Temporary location on the disk of the old, unpacked *image*
         self.old_image_dir = os.path.join(self.tmp_dir, "old")
+        # Temporary location on the disk of the rebase, unpacked *image*
+        self.rebase_image_dir = os.path.join(self.tmp_dir, "rebase")
         # Temporary location on the disk of the new, unpacked, squashed *image*
         self.new_image_dir = os.path.join(self.tmp_dir, "new")
         # Temporary location on the disk of the squashed *layer*
         self.squashed_dir = os.path.join(self.new_image_dir, "squashed")
 
-        for d in self.old_image_dir, self.new_image_dir:
+        for d in self.old_image_dir, self.new_image_dir, self.rebase_image_dir:
             os.makedirs(d)
 
     def _squash_id(self, layer):
@@ -150,8 +155,16 @@ def _before_squashing(self):
         try:
             self.old_image_id = self.docker.inspect_image(self.image)['Id']
         except SquashError:
-            raise SquashError(
-                "Could not get the image ID to squash, please check provided 'image' argument: %s" % self.image)
+            raise SquashError("Could not get the image ID to squash, "
+                              "please check provided 'image' argument: %s" % self.image)
+
+        if self.rebase:
+            # The image id or name of the image to rebase to
+            try:
+                self.rebase = self.docker.inspect_image(self.rebase)['Id']
+            except SquashError:
+                raise SquashError("Could not get the image ID to rebase to, "
+                                  "please check provided 'rebase' argument: %s" % self.rebase)
 
         self.old_image_layers = []
 
@@ -164,32 +177,36 @@ def _before_squashing(self):
         self.log.debug("Old layers: %s", self.old_image_layers)
 
         # By default - squash all layers.
-        if self.from_layer == None:
+        if self.from_layer is None:
             self.from_layer = len(self.old_image_layers)
 
         try:
             number_of_layers = int(self.from_layer)
 
-            self.log.debug(
-                "We detected number of layers as the argument to squash")
+            self.log.debug("We detected number of layers as the argument to squash")
         except ValueError:
             self.log.debug("We detected layer as the argument to squash")
 
             squash_id = self._squash_id(self.from_layer)
 
             if not squash_id:
-                raise SquashError(
-                    "The %s layer could not be found in the %s image" % (self.from_layer, self.image))
+                raise SquashError("The %s layer could not be found in the %s image" % (self.from_layer, self.image))
 
-            number_of_layers = len(self.old_image_layers) - \
-                self.old_image_layers.index(squash_id) - 1
+            number_of_layers = len(self.old_image_layers) - self.old_image_layers.index(squash_id) - 1
 
         self._validate_number_of_layers(number_of_layers)
 
         marker = len(self.old_image_layers) - number_of_layers
 
         self.layers_to_squash = self.old_image_layers[marker:]
-        self.layers_to_move = self.old_image_layers[:marker]
+        if self.rebase:
+            self.layers_to_move = []
+            self._read_layers(self.layers_to_move, self.rebase)
+            self.layers_to_move.reverse()
+        else:
+            self.layers_to_move = self.old_image_layers[:marker]
+
+        self.old_image_squash_marker = marker
 
         self.log.info("Checking if squashing is necessary...")
 
@@ -199,18 +216,18 @@ def _before_squashing(self):
         if len(self.layers_to_squash) == 1:
             raise SquashUnnecessaryError("Single layer marked to squash, no squashing is required")
 
-        self.log.info("Attempting to squash last %s layers...",
-                      number_of_layers)
+        self.log.info("Attempting to squash last %s layers%s...", number_of_layers,
+                      " rebasing on %s" % self.rebase if self.rebase else "")
         self.log.debug("Layers to squash: %s", self.layers_to_squash)
         self.log.debug("Layers to move: %s", self.layers_to_move)
 
         # Fetch the image and unpack it on the fly to the old image directory
         self._save_image(self.old_image_id, self.old_image_dir)
+        if self.rebase:
+            self._save_image(self.rebase, self.rebase_image_dir)
 
         self.size_before = self._dir_size(self.old_image_dir)
 
-        self.log.info("Squashing image '%s'..." % self.image)
-
     def _after_squashing(self):
         self.log.debug("Removing from disk already squashed layers...")
         shutil.rmtree(self.old_image_dir, ignore_errors=True)
@@ -670,7 +687,7 @@ def _squash_layers(self, layers_to_squash, layers_to_move):
 
         # Find all files in layers that we don't squash
         files_in_layers_to_move = self._files_in_layers(
-            layers_to_move, self.old_image_dir)
+            layers_to_move, self.old_image_dir if not self.rebase else self.rebase_image_dir)
 
         with tarfile.open(self.squashed_tar, 'w', format=tarfile.PAX_FORMAT) as squashed_tar:
             to_skip = []
diff --git a/docker_squash/squash.py b/docker_squash/squash.py
index e81960a..9260d07 100644
--- a/docker_squash/squash.py
+++ b/docker_squash/squash.py
@@ -14,7 +14,7 @@
 class Squash(object):
 
     def __init__(self, log, image, docker=None, from_layer=None, tag=None, tmp_dir=None,
-                 output_path=None, load_image=True, development=False, cleanup=False):
+                 output_path=None, load_image=True, development=False, cleanup=False, rebase=None):
         self.log = log
         self.docker = docker
         self.image = image
@@ -25,6 +25,7 @@ def __init__(self, log, image, docker=None, from_layer=None, tag=None, tmp_dir=N
         self.load_image = load_image
         self.development = development
         self.cleanup = cleanup
+        self.rebase = rebase
 
         if not docker:
             self.docker = common.docker_client(self.log)
@@ -48,10 +49,10 @@ def run(self):
 
         if StrictVersion(docker_version['ApiVersion']) >= StrictVersion("1.22"):
             image = V2Image(self.log, self.docker, self.image,
-                            self.from_layer, self.tmp_dir, self.tag)
+                            self.from_layer, self.tmp_dir, self.tag, self.rebase)
         else:
             image = V1Image(self.log, self.docker, self.image,
-                            self.from_layer, self.tmp_dir, self.tag)
+                            self.from_layer, self.tmp_dir, self.tag, self.rebase)
 
         self.log.info("Using %s image format" % image.FORMAT)
 
diff --git a/docker_squash/v1_image.py b/docker_squash/v1_image.py
index c2d3bf4..56945e8 100644
--- a/docker_squash/v1_image.py
+++ b/docker_squash/v1_image.py
@@ -33,7 +33,7 @@ def _squash(self):
         self._write_version_file(self.squashed_dir)
         # Move all the layers that should be untouched
         self._move_layers(self.layers_to_move,
-                          self.old_image_dir, self.new_image_dir)
+                          self.old_image_dir if not self.rebase else self.rebase_image_dir, self.new_image_dir)
 
         config_file = os.path.join(
             self.old_image_dir, self.old_image_id, "json")
diff --git a/docker_squash/v2_image.py b/docker_squash/v2_image.py
index d146852..c35e328 100644
--- a/docker_squash/v2_image.py
+++ b/docker_squash/v2_image.py
@@ -1,10 +1,10 @@
-
 import hashlib
 import json
 import os
 import shutil
-
 from collections import OrderedDict
+from copy import deepcopy
+
 from docker_squash.image import Image
 
 
@@ -22,11 +22,26 @@ def _before_squashing(self):
         self.old_image_config = self._read_json_file(os.path.join(
             self.old_image_dir, self.old_image_manifest['Config']))
 
+        self.rebase_image_manifest = None
+        self.rebase_image_config = None
+        if self.rebase:
+            # Read rebase image manifest file
+            self.rebase_image_manifest = self._read_json_file(
+                os.path.join(self.rebase_image_dir, "manifest.json"))[0]
+
+            # Read rebase image config file
+            self.rebase_image_config = self._read_json_file(os.path.join(
+                self.rebase_image_dir, self.rebase_image_manifest['Config']))
+
         # Read layer paths inside of the tar archive
         # We split it into layers that needs to be squashed
         # and layers that needs to be moved as-is
         self.layer_paths_to_squash, self.layer_paths_to_move = self._read_layer_paths(
-            self.old_image_config, self.old_image_manifest, self.layers_to_move)
+            self.old_image_config, self.old_image_manifest, self.rebase_image_config, self.rebase_image_manifest,
+            self.layers_to_move, self.layers_to_squash)
+
+        self.log.debug("v2: Layer paths to squash: %s", self.layer_paths_to_squash)
+        self.log.debug("v2: Layer paths to move: %s", self.layer_paths_to_move)
 
         if self.layer_paths_to_move:
             self.squash_id = self.layer_paths_to_move[-1]
@@ -64,8 +79,9 @@ def _squash(self):
             shutil.move(self.squashed_dir, os.path.join(
                 self.new_image_dir, layer_path_id))
 
-        manifest = self._generate_manifest_metadata(
-            image_id, self.image_name, self.image_tag, self.old_image_manifest, self.layer_paths_to_move, layer_path_id)
+        manifest = self._generate_manifest_metadata(image_id, self.image_name, self.image_tag,
+                                                    self.old_image_manifest, self.layer_paths_to_move,
+                                                    layer_path_id, rebase_image_manifest=self.rebase_image_manifest)
 
         self._write_manifest_metadata(manifest)
 
@@ -73,7 +89,7 @@ def _squash(self):
 
         # Move all the layers that should be untouched
         self._move_layers(self.layer_paths_to_move,
-                          self.old_image_dir, self.new_image_dir)
+                          self.old_image_dir if not self.rebase else self.rebase_image_dir, self.new_image_dir)
 
         repositories_file = os.path.join(self.new_image_dir, "repositories")
         self._generate_repositories_json(
@@ -104,15 +120,16 @@ def _write_manifest_metadata(self, manifest):
 
         self._write_json_metadata(json_manifest, manifest_file)
 
-    def _generate_manifest_metadata(self, image_id, image_name, image_tag, old_image_manifest, layer_paths_to_move, layer_path_id=None):
+    def _generate_manifest_metadata(self, image_id, image_name, image_tag, old_image_manifest,
+                                    layer_paths_to_move, layer_path_id=None, rebase_image_manifest=None):
         manifest = OrderedDict()
         manifest['Config'] = "%s.json" % image_id
 
         if image_name and image_tag:
             manifest['RepoTags'] = ["%s:%s" % (image_name, image_tag)]
 
-        manifest['Layers'] = old_image_manifest[
-            'Layers'][:len(layer_paths_to_move)]
+        manifest['Layers'] = (old_image_manifest['Layers'][:len(layer_paths_to_move)] if not rebase_image_manifest else
+                              deepcopy(rebase_image_manifest['Layers']))
 
         if layer_path_id:
             manifest['Layers'].append("%s/layer.tar" % layer_path_id)
@@ -127,7 +144,8 @@ def _read_json_file(self, json_file):
         with open(json_file, 'r') as f:
             return json.load(f, object_pairs_hook=OrderedDict)
 
-    def _read_layer_paths(self, old_image_config, old_image_manifest, layers_to_move):
+    def _read_layer_paths(self, old_image_config, old_image_manifest, rebase_image_config, rebase_image_manifest,
+                          layers_to_move, layers_to_squash):
         """
         In case of v2 format, layer id's are not the same as the id's
         used in the exported tar archive to name directories for layers.
@@ -137,11 +155,25 @@ def _read_layer_paths(self, old_image_config, old_image_manifest, layers_to_move
 
         # In manifest.json we do not have listed all layers
         # but only layers that do contain some data.
-        current_manifest_layer = 0
 
         layer_paths_to_move = []
         layer_paths_to_squash = []
 
+        current_manifest_layer = 0
+        if self.rebase:
+            # Iterate over rebase image history, from base image to top layer
+            for i, layer in enumerate(rebase_image_config['history']):
+                # If it's not an empty layer get the id
+                # (directory name) where the layer's data is
+                # stored
+                if not layer.get('empty_layer', False):
+                    layer_id = rebase_image_manifest['Layers'][current_manifest_layer].rsplit('/')[0]
+
+                    # Check if this layer should be moved or squashed
+                    layer_paths_to_move.append(layer_id)
+                    current_manifest_layer += 1
+
+        current_manifest_layer = 0
         # Iterate over image history, from base image to top layer
         for i, layer in enumerate(old_image_config['history']):
 
@@ -149,12 +181,12 @@ def _read_layer_paths(self, old_image_config, old_image_manifest, layers_to_move
             # (directory name) where the layer's data is
             # stored
             if not layer.get('empty_layer', False):
-                layer_id = old_image_manifest['Layers'][
-                    current_manifest_layer].rsplit('/')[0]
+                layer_id = old_image_manifest['Layers'][current_manifest_layer].rsplit('/')[0]
 
                 # Check if this layer should be moved or squashed
-                if len(layers_to_move) > i:
-                    layer_paths_to_move.append(layer_id)
+                if i < self.old_image_squash_marker:
+                    if not self.rebase:
+                        layer_paths_to_move.append(layer_id)
                 else:
                     layer_paths_to_squash.append(layer_id)
 
@@ -163,7 +195,7 @@ def _read_layer_paths(self, old_image_config, old_image_manifest, layers_to_move
         return layer_paths_to_squash, layer_paths_to_move
 
     def _generate_chain_id(self, chain_ids, diff_ids, parent_chain_id):
-        if parent_chain_id == None:
+        if parent_chain_id is None:
             return self._generate_chain_id(chain_ids, diff_ids[1:], diff_ids[0])
 
         chain_ids.append(parent_chain_id)
@@ -188,7 +220,8 @@ def _generate_diff_ids(self):
         diff_ids = []
 
         for path in self.layer_paths_to_move:
-            sha256 = self._compute_sha256(os.path.join(self.old_image_dir, path, "layer.tar"))
+            sha256 = self._compute_sha256(os.path.join(
+                self.old_image_dir if not self.rebase else self.rebase_image_dir, path, "layer.tar"))
             diff_ids.append(sha256)
 
         if self.layer_paths_to_squash:
@@ -303,7 +336,10 @@ def _generate_last_layer_metadata(self, layer_path_id, old_layer_path=None):
     def _generate_image_metadata(self):
         # First - read old image config, we'll update it instead of
         # generating one from scratch
-        metadata = OrderedDict(self.old_image_config)
+        metadata = deepcopy(self.old_image_config)
+        if self.rebase:
+            rebase_metadata = deepcopy(self.rebase_image_config)
+
         # Update image creation date
         metadata['created'] = self.date
 
@@ -311,10 +347,11 @@ def _generate_image_metadata(self):
         metadata.pop("container", None)
 
         # Remove squashed layers from history
-        metadata['history'] = metadata['history'][:len(self.layers_to_move)]
+        metadata['history'] = (metadata['history'][:self.old_image_squash_marker] if not self.rebase else
+                               rebase_metadata['history'])
         # Remove diff_ids for squashed layers
-        metadata['rootfs']['diff_ids'] = metadata['rootfs'][
-            'diff_ids'][:len(self.layer_paths_to_move)]
+        metadata['rootfs']['diff_ids'] = (metadata['rootfs']['diff_ids'][:len(self.layer_paths_to_move)]
+                                          if not self.rebase else rebase_metadata['rootfs']['diff_ids'])
 
         history = {'comment': '', 'created': self.date}
 
diff --git a/tests/test_integ_squash.py b/tests/test_integ_squash.py
index 28fb2fa..0197217 100644
--- a/tests/test_integ_squash.py
+++ b/tests/test_integ_squash.py
@@ -1,19 +1,20 @@
-import unittest
-import mock
-import six
 import codecs
-import os
+import io
 import json
 import logging
+import os
 import shutil
 import tarfile
-import io
-from io import BytesIO
+import unittest
 import uuid
+from io import BytesIO
+
+import mock
+import six
 
-from docker_squash.squash import Squash
 from docker_squash.errors import SquashError, SquashUnnecessaryError
 from docker_squash.lib import common
+from docker_squash.squash import Squash
 
 if not six.PY3:
     import docker_squash.lib.xtarfile
@@ -21,7 +22,7 @@
 class ImageHelper(object):
     @staticmethod
     def top_layer_path(tar):
-        #tar_object.seek(0)
+        # tar_object.seek(0)
         reader = codecs.getreader("utf-8")
 
         if 'repositories' in tar.getnames():
@@ -34,8 +35,8 @@ def top_layer_path(tar):
             manifest = json.load(reader(tar.extractfile(manifest_member)))
             return manifest[0]["Layers"][-1].split("/")[0]
 
-class IntegSquash(unittest.TestCase):
 
+class IntegSquash(unittest.TestCase):
     BUSYBOX_IMAGE = "busybox:1.24"
 
     log = logging.getLogger()
@@ -99,7 +100,8 @@ def _save_image(self):
 
     class SquashedImage(object):
 
-        def __init__(self, image, number_of_layers=None, output_path=None, load_image=True, numeric=False, tmp_dir=None, log=None, development=False, tag=True):
+        def __init__(self, image, number_of_layers=None, output_path=None, load_image=True, numeric=False, tmp_dir=None,
+                     log=None, development=False, tag=True, rebase=None):
             self.image = image
             self.number_of_layers = number_of_layers
             self.docker = TestIntegSquash.docker
@@ -113,17 +115,15 @@ def __init__(self, image, number_of_layers=None, output_path=None, load_image=Tr
             self.numeric = numeric
             self.tmp_dir = tmp_dir
             self.development = development
+            self.rebase = rebase
 
         def __enter__(self):
             from_layer = self.number_of_layers
 
-            if self.number_of_layers and not self.numeric:
-                from_layer = self.docker.history(
-                    self.image.tag)[self.number_of_layers]['Id']
-
             squash = Squash(
                 self.log, self.image.tag, self.docker, tag=self.tag, from_layer=from_layer,
-                output_path=self.output_path, load_image=self.load_image, tmp_dir=self.tmp_dir, development=self.development)
+                output_path=self.output_path, load_image=self.load_image, tmp_dir=self.tmp_dir,
+                development=self.development, rebase=self.rebase)
 
             self.image_id = squash.run()
 
@@ -184,7 +184,7 @@ def assertFileIsNotHardLink(self, name):
             with tarfile.open(fileobj=self.squashed_layer, mode='r') as tar:
                 member = tar.getmember(name)
                 assert member.islnk(
-                ) == False, "File '%s' should not be a hard link, but it is" % name
+                ) is False, "File '%s' should not be a hard link, but it is" % name
 
     class Container(object):
 
@@ -215,6 +215,7 @@ def assertFileDoesNotExist(self, name):
                 assert name not in tar.getnames(
                 ), "File %s was found in the container files: %s" % (name, tar.getnames())
 
+
 class TestIntegSquash(IntegSquash):
 
     def test_all_files_should_be_in_squashed_layer(self):
@@ -848,7 +849,6 @@ def test_should_not_skip_sym_link(self):
 
         with self.Image(dockerfile) as image:
             with self.SquashedImage(image, 2, numeric=True) as squashed_image:
-
                 with self.Container(squashed_image) as container:
                     container.assertFileExists('dir')
                     container.assertFileExists('dir/a')
@@ -872,7 +872,6 @@ def test_should_not_skip_hard_link(self):
 
         with self.Image(dockerfile) as image:
             with self.SquashedImage(image, 2, numeric=True) as squashed_image:
-
                 with self.Container(squashed_image) as container:
                     container.assertFileExists('dir')
                     container.assertFileExists('dir/a')
@@ -937,7 +936,8 @@ def test_should_not_add_duplicate_files(self):
                     container.assertFileExists('data-template/etc/systemd/system/default.target.wants')
                     container.assertFileExists('data-template/etc/systemd/system/default.target')
                     container.assertFileExists('data-template/etc/systemd/system/multi-user.target.wants')
-                    container.assertFileExists('data-template/etc/systemd/system/container-ipa.target.wants/ipa-server-configure-first.service')
+                    container.assertFileExists(
+                        'data-template/etc/systemd/system/container-ipa.target.wants/ipa-server-configure-first.service')
                     container.assertFileExists('etc/systemd/system')
 
 
@@ -980,7 +980,6 @@ def test_should_not_squash_single_layer(self):
 
     def test_should_squash_2_layers(self):
         with self.SquashedImage(NumericValues.image, 2, numeric=True) as squashed_image:
-
             i_h = NumericValues.image.history[0]
             s_h = squashed_image.history[0]
 
@@ -1012,5 +1011,41 @@ def test_should_squash_4_layers(self):
             self.assertEqual(
                 len(squashed_image.layers), len(NumericValues.image.layers) - 3)
 
+
+class RebaseTests(IntegSquash):
+    def test_rebase(self):
+        dockerfile_base = '''
+        FROM %s
+        RUN touch /layer_that_stays
+        ''' % TestIntegSquash.BUSYBOX_IMAGE
+
+        with self.Image(dockerfile_base) as base_image:
+            dockerfile_dev_base = '''
+            FROM %s
+            RUN touch /somefile_layer1
+            RUN touch /somefile_layer2
+            RUN touch /somefile_layer3
+            ''' % base_image.tag
+
+            with self.Image(dockerfile_dev_base) as dev_base_image:
+                dockerfile_final = '''
+                FROM %s
+                RUN touch /somefile_layer4
+                RUN touch /somefile_layer5
+                ''' % dev_base_image.tag
+
+                with self.Image(dockerfile_final) as final_image:
+                    with self.SquashedImage(final_image, rebase=base_image.tag,
+                                            number_of_layers=dev_base_image.tag) as squashed_image:
+                        with self.Container(squashed_image) as container:
+                            container.assertFileDoesNotExist('somefile_layer1')
+                            container.assertFileDoesNotExist('somefile_layer2')
+                            container.assertFileDoesNotExist('somefile_layer3')
+                            container.assertFileExists('somefile_layer4')
+                            container.assertFileExists('somefile_layer5')
+                            container.assertFileExists('bin/sh')
+                            container.assertFileExists('layer_that_stays')
+
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_unit_squash.py b/tests/test_unit_squash.py
index 0fa11fc..90b0740 100644
--- a/tests/test_unit_squash.py
+++ b/tests/test_unit_squash.py
@@ -19,7 +19,7 @@ def test_handle_case_when_no_image_is_provided(self):
         squash = Squash(self.log, None, self.docker_client)
         with self.assertRaises(SquashError) as cm:
             squash.run()
-        self.assertEquals(
+        self.assertEqual(
             str(cm.exception), "Image is not provided")
 
     def test_exit_if_no_output_path_provided_and_loading_is_disabled_too(self):
diff --git a/tests/test_unit_v2_image.py b/tests/test_unit_v2_image.py
index 251df8c..f391b98 100644
--- a/tests/test_unit_v2_image.py
+++ b/tests/test_unit_v2_image.py
@@ -43,13 +43,14 @@ def setUp(self):
         self.image = "whatever"
         self.image = V2Image(self.log, self.docker_client, self.image, None)
 
-    def test_generate_manifest(self):
+    def test_generate_manifest_no_rebase(self):
         old_image_manifest = {'Layers': [
             "layer_a/layer.tar", "layer_b/layer.tar", "layer_c/layer.tar"]}
         layer_paths_to_move = ["layer_a", "layer_b"]
 
         metadata = self.image._generate_manifest_metadata(
-            "this_is_image_id", "image", "squashed", old_image_manifest, layer_paths_to_move, "this_is_layer_path_id")
+            "this_is_image_id", "image", "squashed", old_image_manifest, layer_paths_to_move,
+            "this_is_layer_path_id")
 
         self.assertEqual(len(metadata), 1)
 
@@ -61,6 +62,27 @@ def test_generate_manifest(self):
         self.assertEqual(metadata['Layers'], [
                          "layer_a/layer.tar", "layer_b/layer.tar", "this_is_layer_path_id/layer.tar"])
 
+    def test_generate_manifest_rebase(self):
+        old_image_manifest = {'Layers': [
+            "layer_a/layer.tar", "layer_b/layer.tar", "layer_c/layer.tar"]}
+        rebase_image_manifest = {'Layers': [
+            "layer_x/layer.tar", "layer_y/layer.tar", "layer_z/layer.tar"]}
+        layer_paths_to_move = ["layer_x", "layer_y", "layer_z"]
+
+        metadata = self.image._generate_manifest_metadata(
+            "this_is_image_id", "image", "squashed", old_image_manifest, layer_paths_to_move,
+            "this_is_layer_path_id", rebase_image_manifest=rebase_image_manifest)
+
+        self.assertEqual(len(metadata), 1)
+
+        metadata = metadata[0]
+
+        self.assertEqual(type(metadata), OrderedDict)
+        self.assertEqual(metadata['Config'], 'this_is_image_id.json')
+        self.assertEqual(metadata['RepoTags'], ['image:squashed'])
+        self.assertEqual(metadata['Layers'], ["layer_x/layer.tar", "layer_y/layer.tar", "layer_z/layer.tar",
+                                              "this_is_layer_path_id/layer.tar"])
+
     def test_generate_image_metadata_without_any_layers_to_squash(self):
         self.image.old_image_dir = "/tmp/old"
         self.image.squash_id = "squash_id"
@@ -71,6 +93,7 @@ def test_generate_image_metadata_without_any_layers_to_squash(self):
         # We want to move 2 layers with content
         self.image.layer_paths_to_move = ["layer_path_1", "layer_path_2"]
         self.image.layer_paths_to_squash = []
+        self.image.old_image_squash_marker = 3
         # Image that contains:
         # - 4 layers
         # - 3 layers that have content
@@ -97,6 +120,7 @@ def test_generate_image_metadata(self):
         # We want to move 2 layers with content
         self.image.layer_paths_to_move = ["layer_path_1", "layer_path_2"]
         self.image.layer_paths_to_squash = ["layer_path_3", "layer_path_4"]
+        self.image.old_image_squash_marker = 3
         # Image that contains:
         # - 4 layers
         # - 3 layers that have content
diff --git a/tox.ini b/tox.ini
index 90c8790..f147817 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = py27,py34,py35
+envlist = py27,py34,py35,py36
 
 [testenv]
 passenv=CI