From 0464a3ebfca9219dc01a3fe47e06409c4d88af85 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Mon, 14 Aug 2017 09:43:18 +0200
Subject: [PATCH] tools: implement qvm-create --root-move-from/--root-copy-from

---
 qubesadmin/tests/tools/qvm_create.py | 67 ++++++++++++++++++++++++++++
 qubesadmin/tools/qvm_create.py       | 18 +++++++-
 2 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/qubesadmin/tests/tools/qvm_create.py b/qubesadmin/tests/tools/qvm_create.py
index bf934e4..1723e26 100644
--- a/qubesadmin/tests/tools/qvm_create.py
+++ b/qubesadmin/tests/tools/qvm_create.py
@@ -17,6 +17,9 @@
 #
 # You should have received a copy of the GNU Lesser General Public License along
 # with this program; if not, see <http://www.gnu.org/licenses/>.
+import os
+import tempfile
+
 import qubesadmin.tests
 import qubesadmin.tests.tools
 import qubesadmin.tools.qvm_create
@@ -89,3 +92,67 @@ class TC_00_qvm_create(qubesadmin.tests.QubesTestCase):
             'private=some-pool', '--pool', 'volatile=other-pool', 'new-vm'],
             app=self.app)
         self.assertAllCalled()
+
+    def test_005_root_copy_from(self):
+        with tempfile.NamedTemporaryFile() as root_file:
+            root_file.file.write(b'root data')
+            root_file.file.flush()
+            self.app.expected_calls[('dom0', 'admin.vm.Create.AppVM',
+                None, b'name=new-vm label=red')] = b'0\x00'
+            self.app.expected_calls[('dom0', 'admin.label.List', None, None)] = \
+                b'0\x00red\nblue\n'
+            self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
+                b'0\x00new-vm class=AppVM state=Halted\n'
+            self.app.expected_calls[
+                ('new-vm', 'admin.vm.volume.List', None, None)] = \
+                b'0\x00root\nprivate\nvolatile\nkernel\n'
+            self.app.expected_calls[
+                ('new-vm', 'admin.vm.volume.Import', 'root', b'root data')] = \
+                b'0\0'
+            qubesadmin.tools.qvm_create.main(['-l', 'red',
+                '--root-copy-from=' + root_file.name, 'new-vm'],
+                app=self.app)
+            self.assertAllCalled()
+            self.assertTrue(os.path.exists(root_file.name))
+
+    def test_006_root_move_from(self):
+        with tempfile.NamedTemporaryFile(delete=False) as root_file:
+            root_file.file.write(b'root data')
+            root_file.file.flush()
+            self.app.expected_calls[('dom0', 'admin.vm.Create.AppVM',
+                None, b'name=new-vm label=red')] = b'0\x00'
+            self.app.expected_calls[('dom0', 'admin.label.List', None, None)] = \
+                b'0\x00red\nblue\n'
+            self.app.expected_calls[('dom0', 'admin.vm.List', None, None)] = \
+                b'0\x00new-vm class=AppVM state=Halted\n'
+            self.app.expected_calls[
+                ('new-vm', 'admin.vm.volume.List', None, None)] = \
+                b'0\x00root\nprivate\nvolatile\nkernel\n'
+            self.app.expected_calls[
+                ('new-vm', 'admin.vm.volume.Import', 'root', b'root data')] = \
+                b'0\0'
+            qubesadmin.tools.qvm_create.main(['-l', 'red',
+                '--root-move-from=' + root_file.name, 'new-vm'],
+                app=self.app)
+            self.assertAllCalled()
+            self.assertFalse(os.path.exists(root_file.name))
+
+    def test_007_root_move_copy_both(self):
+        with tempfile.NamedTemporaryFile() as root_file:
+            root_file.file.write(b'root data')
+            root_file.file.flush()
+            with self.assertRaises(SystemExit):
+                qubesadmin.tools.qvm_create.main(['-l', 'red',
+                    '--root-copy-from=' + root_file.name,
+                    '--root-move-from=' + root_file.name,
+                    'new-vm'],
+                    app=self.app)
+            self.assertAllCalled()
+            self.assertTrue(os.path.exists(root_file.name))
+
+    def test_008_root_invalid_path(self):
+        with self.assertRaises(SystemExit):
+            qubesadmin.tools.qvm_create.main(['-l', 'red',
+                '--root-copy-from=/invalid', 'new-vm'],
+                app=self.app)
+        self.assertAllCalled()
diff --git a/qubesadmin/tools/qvm_create.py b/qubesadmin/tools/qvm_create.py
index 91bf9c7..10bebcb 100644
--- a/qubesadmin/tools/qvm_create.py
+++ b/qubesadmin/tools/qvm_create.py
@@ -30,6 +30,7 @@
 from __future__ import print_function
 
 import argparse
+import os
 import sys
 
 import qubesadmin
@@ -108,9 +109,10 @@ def main(args=None, app=None):
     if 'name' not in args.properties:
         parser.error('VMNAME is mandatory')
 
-    if args.root_copy_from or args.root_move_from:
+    root_source_path = args.root_copy_from or args.root_move_from
+    if root_source_path and not os.path.exists(root_source_path):
         parser.error(
-            '--root-copy-from and --root-move-from not implemented yet')
+            'File pointed by --root-copy-from/--root-move-from does not exist')
 
     try:
         args.app.get_label(args.properties['label'])
@@ -145,6 +147,18 @@ def main(args=None, app=None):
                 format(prop, e))
             retcode = 2
 
+    if root_source_path:
+        try:
+            with open(root_source_path, 'rb') as root_file:
+                vm.volumes['root'].import_data(root_file)
+            if args.root_move_from:
+                os.unlink(root_source_path)
+        except qubesadmin.exc.QubesException as e:
+            args.app.log.error(
+                'Error importing root volume (but VM created): {}'.
+                format(e))
+            retcode = 3
+
     return retcode