diff --git a/distrobuilder/main_lxc.go b/distrobuilder/main_lxc.go
index 40a1cd9..8690ef7 100644
--- a/distrobuilder/main_lxc.go
+++ b/distrobuilder/main_lxc.go
@@ -58,7 +58,8 @@ func (c *cmdLXC) run(cmd *cobra.Command, args []string) error {
 			continue
 		}
 
-		err := generator.CreateLXCData(c.global.flagCacheDir, file.Path, img)
+		err := generator.CreateLXCData(c.global.flagCacheDir, c.global.sourceDir,
+			file.Path, img)
 		if err != nil {
 			continue
 		}
@@ -70,7 +71,7 @@ func (c *cmdLXC) run(cmd *cobra.Command, args []string) error {
 	}
 
 	// Clean up the chroot by restoring the orginal files.
-	err = generators.RestoreFiles(c.global.flagCacheDir)
+	err = generators.RestoreFiles(c.global.flagCacheDir, c.global.sourceDir)
 	if err != nil {
 		return fmt.Errorf("Failed to restore cached files: %s", err)
 	}
diff --git a/distrobuilder/main_lxd.go b/distrobuilder/main_lxd.go
index d109a1c..d695b9a 100644
--- a/distrobuilder/main_lxd.go
+++ b/distrobuilder/main_lxd.go
@@ -76,7 +76,8 @@ func (c *cmdLXD) run(cmd *cobra.Command, args []string) error {
 			return fmt.Errorf("Unknown generator '%s'", file.Generator)
 		}
 
-		err := generator.CreateLXDData(c.global.sourceDir, file.Path, img)
+		err := generator.CreateLXDData(c.global.flagCacheDir, c.global.sourceDir,
+			file.Path, img)
 		if err != nil {
 			return fmt.Errorf("Failed to create LXD data: %s", err)
 		}
diff --git a/generators/generators.go b/generators/generators.go
index 19c55dd..c126f14 100644
--- a/generators/generators.go
+++ b/generators/generators.go
@@ -13,8 +13,8 @@ import (
 
 // Generator interface.
 type Generator interface {
-	CreateLXCData(string, string, *image.LXCImage) error
-	CreateLXDData(string, string, *image.LXDImage) error
+	CreateLXCData(string, string, string, *image.LXCImage) error
+	CreateLXDData(string, string, string, *image.LXDImage) error
 }
 
 // Get returns a Generator.
@@ -30,19 +30,19 @@ func Get(generator string) Generator {
 }
 
 // StoreFile caches a file which can be restored with the RestoreFiles function.
-func StoreFile(cacheDir, path string) error {
+func StoreFile(cacheDir, sourceDir, path string) error {
 	// create temporary directory containing old files
 	err := os.MkdirAll(filepath.Join(cacheDir, "tmp", p.Dir(path)), 0755)
 	if err != nil {
 		return err
 	}
 
-	return shared.FileCopy(filepath.Join(cacheDir, "rootfs", path),
+	return shared.FileCopy(filepath.Join(sourceDir, path),
 		filepath.Join(cacheDir, "tmp", path))
 }
 
 // RestoreFiles restores original files which were cached by StoreFile.
-func RestoreFiles(cacheDir string) error {
+func RestoreFiles(cacheDir, sourceDir string) error {
 	f := func(path string, info os.FileInfo, err error) error {
 		if info.IsDir() {
 			// We don't care about directories. They should be present so there's
@@ -50,7 +50,7 @@ func RestoreFiles(cacheDir string) error {
 			return nil
 		}
 
-		return os.Rename(path, filepath.Join(cacheDir, "rootfs",
+		return shared.FileCopy(path, filepath.Join(sourceDir,
 			strings.TrimPrefix(path, filepath.Join(cacheDir, "tmp"))))
 	}
 
diff --git a/generators/hostname.go b/generators/hostname.go
index f0b1206..ff47a00 100644
--- a/generators/hostname.go
+++ b/generators/hostname.go
@@ -14,17 +14,15 @@ import (
 type HostnameGenerator struct{}
 
 // CreateLXCData creates a hostname template.
-func (g HostnameGenerator) CreateLXCData(cacheDir, path string, img *image.LXCImage) error {
-	rootfs := filepath.Join(cacheDir, "rootfs")
-
+func (g HostnameGenerator) CreateLXCData(cacheDir, sourceDir, path string, img *image.LXCImage) error {
 	// Store original file
-	err := StoreFile(cacheDir, path)
+	err := StoreFile(cacheDir, sourceDir, path)
 	if err != nil {
 		return err
 	}
 
 	// Create new hostname file
-	file, err := os.Create(filepath.Join(rootfs, path))
+	file, err := os.Create(filepath.Join(sourceDir, path))
 	if err != nil {
 		return err
 	}
@@ -41,7 +39,7 @@ func (g HostnameGenerator) CreateLXCData(cacheDir, path string, img *image.LXCIm
 }
 
 // CreateLXDData creates a hostname template.
-func (g HostnameGenerator) CreateLXDData(cacheDir, path string, img *image.LXDImage) error {
+func (g HostnameGenerator) CreateLXDData(cacheDir, sourceDir, path string, img *image.LXDImage) error {
 	templateDir := filepath.Join(cacheDir, "templates")
 
 	err := os.MkdirAll(templateDir, 0755)
diff --git a/generators/hosts.go b/generators/hosts.go
index 489f0f5..177ced8 100644
--- a/generators/hosts.go
+++ b/generators/hosts.go
@@ -13,16 +13,14 @@ import (
 type HostsGenerator struct{}
 
 // CreateLXCData creates a LXC specific entry in the hosts file.
-func (g HostsGenerator) CreateLXCData(cacheDir, path string, img *image.LXCImage) error {
-	rootfs := filepath.Join(cacheDir, "rootfs")
-
+func (g HostsGenerator) CreateLXCData(cacheDir, sourceDir, path string, img *image.LXCImage) error {
 	// Store original file
-	err := StoreFile(cacheDir, path)
+	err := StoreFile(cacheDir, sourceDir, path)
 	if err != nil {
 		return err
 	}
 
-	file, err := os.OpenFile(filepath.Join(rootfs, path),
+	file, err := os.OpenFile(filepath.Join(sourceDir, path),
 		os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
 	if err != nil {
 		return err
@@ -37,7 +35,7 @@ func (g HostsGenerator) CreateLXCData(cacheDir, path string, img *image.LXCImage
 }
 
 // CreateLXDData creates a hosts template.
-func (g HostsGenerator) CreateLXDData(cacheDir, path string, img *image.LXDImage) error {
+func (g HostsGenerator) CreateLXDData(cacheDir, sourceDir, path string, img *image.LXDImage) error {
 	templateDir := filepath.Join(cacheDir, "templates")
 
 	// Create templates path
@@ -53,7 +51,7 @@ func (g HostsGenerator) CreateLXDData(cacheDir, path string, img *image.LXDImage
 	}
 	defer file.Close()
 
-	hostsFile, err := os.Open(filepath.Join(cacheDir, "rootfs", path))
+	hostsFile, err := os.Open(filepath.Join(sourceDir, path))
 	if err != nil {
 		return err
 	}