Create disk images in multiple formats

Refactor some code in create_disk_images.oy into a »prepare_disks«
function that creates disk images in the given format for one task.
Store paths to images in multiple formats in the database and update
the front-end.

4 files changed, 159 insertions, 140 deletions

diff --git a/kpov_judge/create_disk_images.py b/kpov_judge/create_disk_images.py
index ce79f85..381ba25 100755
--- a/kpov_judge/create_disk_images.py
+++ b/kpov_judge/create_disk_images.py
@@ -23,27 +23,30 @@ def get_prepare_disks(db, class_id, task_id):
     exec(compile(prepare_disks_source, 'prepare_disks.py', 'exec'), globals(), d)
     return d['prepare_disks']
 
-def create_snapshot(class_id, task_id, student_id, disk_name, overwrite=True):
+def create_snapshot(class_id, task_id, student_id, disk_name, fmt='vmdk', overwrite=True):
     # add a hash to filename to allow multiple students using the same directory
     snap_hash = hashlib.sha1((disk_name+class_id+task_id+student_id).encode()).hexdigest()[:4]
     snap = '{}-{}-{}.{}'.format(
-        task_id, disk_name, snap_hash, settings.STUDENT_DISK_FORMAT)
+        task_id, disk_name, snap_hash, fmt)
     backing = []
 
-    template = disk_name + '.' + settings.STUDENT_DISK_FORMAT
+    template = disk_name + '.' + fmt
     task_dir = os.path.join(student_id, class_id, task_id)
     task_path = os.path.join(settings.STUDENT_DISK_PATH, task_dir)
 
     if not os.path.exists(os.path.join(task_path)) or overwrite:
-        if settings.STUDENT_DISK_COW:
-            # don’t use backing files, just copy the template
-            # (requires a cow-capable filesystem)
-            subprocess.call(['cp', '--reflink=always', template, snap])
+        # ensure task dir exists
+        os.makedirs(task_path, exist_ok=True)
 
-        else:
-            # create task dir
-            os.makedirs(task_path, exist_ok=True)
+        if fmt == 'vdi':
+            # don’t use backing files, just copy the template
+            os.chdir(task_path)
+            if settings.STUDENT_DISK_COW:
+                subprocess.call(['cp', '--reflink=always', os.path.join(settings.DISK_TEMPLATE_PATH, template), snap])
+            else:
+                subprocess.call(['cp', os.path.join(settings.DISK_TEMPLATE_PATH, template), snap])
 
+        elif fmt == 'vmdk':
             # qemu-img create stores backing-file path as given, so link all
             # backing images to task directory where target image will be
             # generated
@@ -60,11 +63,93 @@ def create_snapshot(class_id, task_id, student_id, disk_name, overwrite=True):
             # make overlay image
             os.chdir(task_path)
             subprocess.call(['qemu-img', 'create',
-                '-f', settings.STUDENT_DISK_FORMAT,
+                '-f', fmt,
                 '-b', template, snap])
 
     return task_dir, snap, backing
 
+def prepare_task_disks(class_id, task_id, student_id, fmt, computers):
+    disks = collections.defaultdict(dict)
+    templates = collections.defaultdict(dict)
+    for computer in computers:
+        lock_fp.write('creating computer ' + computer['name'] + '\n')
+        if not computer['disks']:
+            continue
+
+        manual_disks = []
+        try_automount = False
+
+        g = guestfs.GuestFS()
+        for disk in computer['disks']:
+            lock_fp.write("register " + disk['name'] + '\n')
+            task_dir, snap, backing = create_snapshot(class_id, task_id, student_id, disk['name'], fmt=fmt)
+            snap_file = os.path.join(settings.STUDENT_DISK_PATH, task_dir, snap)
+            if 'options' in disk:
+                g.add_drive_opts(snap_file, **(disk['options']))
+            else:
+                g.add_drive(snap_file)
+            if 'parts' in disk:
+                for p in disk['parts']:
+                    lock_fp.write("part {}: {}\n".format(
+                        settings.GUESTFS_DEV_PREFIX + p['dev'], p['path']))
+                    manual_disks.append(
+                        (settings.GUESTFS_DEV_PREFIX + p['dev'], p['path'], p.get('options', None)))
+            else:
+                try_automount = True
+
+            templates[disk['name']] = g
+            lock_fp.write("  templates[{}] = {}\n".format(disk['name'], disk))
+
+            # add disk or update existing record with new format
+            disks[computer['name']][disk['name']] = [snap] + backing
+
+        g.launch()
+        mounted = set()
+        if try_automount:
+            roots = g.inspect_os()
+            for root in roots:
+                mps = g.inspect_get_mountpoints(root)
+                lock_fp.write('detected: ' + str(mps) + '\n')
+                for mountpoint, device in sorted(mps):
+                    if mountpoint not in mounted:
+                        try:
+                            g.mount(device, mountpoint, )
+                            lock_fp.write( 'mounted ' + device + ' on ' + mountpoint + '\n')
+                        except RuntimeError as msg:
+                            lock_fp.write( "%s (ignored)\n" % msg)
+                        mounted.add(mountpoint)
+
+        for device, mountpoint, opts in manual_disks:
+            try:
+                if opts is not None:
+                    g.mount_options(opts, device, mountpoint)
+                else:
+                    g.mount(device, mountpoint)
+                lock_fp.write('manually mounted ' + device + " on " + mountpoint + '\n')
+            except RuntimeError as msg:
+                lock_fp.write( "%s (ignored)\n" % msg)
+
+    lock_fp.write("preparing disks\n")
+    global_params = {
+        'task_name': task_id,
+        'class_id': class_id,
+        'username': student_id
+    }
+    if 'TASK_URL' in vars(settings):
+        global_params['task_url'] = settings.TASK_URL + '/' + class_id + '/'
+
+    task_params = db.task_params.find_one({'class_id': class_id, 'task_id': task_id, 'student_id': student_id})['params']
+    prepare_disks = get_prepare_disks(db, class_id, task_id)
+    prepare_disks(templates, task_params, global_params)
+
+    # pospravi za seboj.
+    lock_fp.write("unmounting\n")
+    for g in set(templates.values()):
+        g.umount_all()
+        g.close()
+
+    return disks
+
 if __name__ == '__main__':
     if len(sys.argv) != 1:
         print("Usage: {0}")
@@ -77,119 +162,37 @@ if __name__ == '__main__':
         all_computers[(computer['class_id'], computer['task_id'], computer['student_id'])] += [computer]
 
     for (class_id, task_id, student_id), computers in all_computers.items():
-        # TODO check why we iterate over student_computers twice
-        l = db.student_computers.find_one({'class_id': class_id, 'task_id': task_id, 'student_id': student_id})
-        if l is None:
+        if db.student_computers.find_one({'class_id': class_id, 'task_id': task_id, 'student_id': student_id}) is None:
             continue
 
-        print("Creating {}/{} for {}".format(class_id, task_id, student_id))
-
         lock_file = os.path.join(settings.STUDENT_LOCKFILE_PATH, 
             '{0}-{1}-{2}.lock'.format(student_id, class_id, task_id))
-        lock_fp = open(lock_file, 'w')
-        try:
-            fcntl.lockf(lock_fp, fcntl.LOCK_EX | fcntl.LOCK_NB)
-        except IOError:
-            continue
-
-        task_params = db.task_params.find_one({'class_id': class_id, 'task_id': task_id, 'student_id': student_id})['params']
-        prepare_disks = get_prepare_disks(db, class_id, task_id)
-
-        # tule odpri, ustvari snapshote za vajo
-        templates = {}
-        all_disks = collections.defaultdict(list)
-        parts = {}
-        for computer in computers:
-            lock_fp.write('creating computer ' + computer['name'] + '\n')
-            if len(computer['disks']) == 0:
+        with open(lock_file, 'w') as lock_fp:
+            try:
+                fcntl.lockf(lock_fp, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            except IOError:
                 continue
 
-            manual_disks = []
-            try_automount = False
-
-            g = guestfs.GuestFS()
-            for disk in computer['disks']:
-                lock_fp.write("register " + disk['name'] + '\n')
-                task_dir, snap, backing = create_snapshot(class_id, task_id, student_id, disk['name'])
-                snap_file = os.path.join(settings.STUDENT_DISK_PATH, task_dir, snap)
-                if 'options' in disk:
-                    g.add_drive_opts(snap_file, **(disk['options']))
-                else:
-                    g.add_drive(snap_file)
-                if 'parts' in disk:
-                    for p in disk['parts']:
-                        lock_fp.write("part {}: {}\n".format(
-                            settings.GUESTFS_DEV_PREFIX + p['dev'], p['path']))
-                        manual_disks.append(
-                            (settings.GUESTFS_DEV_PREFIX + p['dev'], p['path'], p.get('options', None)))
-                else:
-                    try_automount = True
-
-                templates[disk['name']] = g
-                lock_fp.write("  templates[{}] = {}\n".format(disk['name'], disk))
-
-                all_disks[computer['name']] += [{
-                    'name': disk['name'],
-                    'file': snap,
-                    'backing': backing,
-                }]
-
-            g.launch()
-            mounted = set()
-            if try_automount:
-                roots = g.inspect_os()
-                for root in roots:
-                    mps = g.inspect_get_mountpoints(root)
-                    lock_fp.write('detected: ' + str(mps) + '\n')
-                    for mountpoint, device in sorted(mps):
-                        if mountpoint not in mounted:
-                            try:
-                                g.mount(device, mountpoint, )
-                                lock_fp.write( 'mounted ' + device + ' on ' + mountpoint + '\n')
-                            except RuntimeError as msg:
-                                lock_fp.write( "%s (ignored)\n" % msg)
-                            mounted.add(mountpoint)
-
-            for device, mountpoint, opts in manual_disks:
-                try:
-                    if opts is not None:
-                        g.mount_options(opts, device, mountpoint)
-                    else:
-                        g.mount(device, mountpoint)
-                    lock_fp.write('manually mounted ' + device + " on " + mountpoint + '\n')
-                except RuntimeError as msg:
-                    lock_fp.write( "%s (ignored)\n" % msg)
-
-        lock_fp.write("preparing disks\n")
-        global_params = {
-            'task_name': task_id,
-            'class_id': class_id,
-            'username': student_id
-        }
-        if 'TASK_URL' in vars(settings):
-            global_params['task_url'] = settings.TASK_URL + '/' + class_id + '/'
-
-        prepare_disks(templates, task_params, global_params)
-
-        # pospravi za seboj.
-        lock_fp.write("unmounting\n")
-        for g in set(templates.values()):
-            g.umount_all()
-            g.close()
-
-        lock_fp.write("saving URLs\n")
-        for computer in computers:
-            comp_name = computer['name']
-            disks = all_disks[comp_name]
-            lock_fp.write('urls: '+ str(disks) + '\n')
-            l = db.student_computers.update({
-                'disk_urls': {'$exists': False}, 
-                'student_id': student_id, 
-                'task_id': task_id,
-                'class_id': class_id,
-                'name': comp_name},
-                {'$set': { 'disk_urls': disks }
-            })
-
-        os.unlink(lock_file)
-        lock_fp.close()
+            print("Creating {}/{} for {}".format(class_id, task_id, student_id))
+            all_disks = collections.defaultdict(dict)
+            for fmt in settings.STUDENT_DISK_FORMATS:
+                for computer, disks in prepare_task_disks(class_id, task_id, student_id, fmt, computers).items():
+                    for disk, urls in disks.items():
+                        d = all_disks[computer].setdefault(disk, {'formats': []})
+                        d['formats'] += [fmt]
+                        d[fmt] = urls
+
+            lock_fp.write("saving URLs\n")
+            for computer in computers:
+                comp_name = computer['name']
+                disks = all_disks[comp_name]
+                lock_fp.write('urls: '+ str(disks) + '\n')
+                db.student_computers.update({
+                        'disk_urls': {'$exists': False},
+                        'student_id': student_id,
+                        'task_id': task_id,
+                        'class_id': class_id,
+                        'name': comp_name},
+                    {'$set': { 'disk_urls': disks }})
+
+            os.unlink(lock_file)
diff --git a/kpov_judge/settings-example.py b/kpov_judge/settings-example.py
index 00003f7..34e70c3 100644
--- a/kpov_judge/settings-example.py
+++ b/kpov_judge/settings-example.py
@@ -4,7 +4,7 @@ DB_URI='mongodb://kpov:SomeSuperSecretPassword@kpov.fri.uni-lj.si:27017/kpov'
 # SNAPSHOTS_ROOT='/home/andrejtolic/kpov/snapshots/'
 DISK_TEMPLATE_PATH = '/home/kpov/disks/templates'
 STUDENT_DISK_PATH='/home/kpov/disks/students'
-STUDENT_DISK_FORMAT='vdi'
+STUDENT_DISK_FORMAT=['vdi', 'vmdk']
 STUDENT_DISK_COW=False # dokler Andrej ne popravi...
 STUDENT_DISK_URL='https://my_html_server.com/kpov-disks/'
 STUDENT_LOCKFILE_PATH='/home/kpov/disks/lockfiles'
diff --git a/kpov_judge/web/kpov_judge/kpov_judge.py b/kpov_judge/web/kpov_judge/kpov_judge.py
index c0bf885..9e4e24a 100755
--- a/kpov_judge/web/kpov_judge/kpov_judge.py
+++ b/kpov_judge/web/kpov_judge/kpov_judge.py
@@ -255,11 +255,17 @@ def task_greeting(class_id, task_id, lang):
                     }]
         except Exception as e:
             instructions = str(e)
+
     computer_list = list(db.student_computers.find({'class_id': class_id, 'task_id': task_id, 'student_id': student_id}))
-    backing_images = set()
+
+    backing_files = collections.defaultdict(set)
     for computer in computer_list:
-        for disk_url in computer.get('disk_urls', []):
-            backing_images |= set(disk_url['backing'])
+        if 'disk_urls' not in computer:
+            continue
+        for name, disk in computer['disk_urls'].items():
+            for fmt in disk['formats']:
+                backing_files[fmt] |= set(disk[fmt][1:])
+
     if request.args.get('narediStack', 'false') == 'true':
       #db.student_tasks.update({'task_id': task_id, 'student_id': student_id}, {'$set': {'create_openstack': True}}, upsert = True)
         openstackCreated = False # Spremeni na True, ko odkomentiras zgornjo vrstico.
@@ -285,8 +291,8 @@ def task_greeting(class_id, task_id, lang):
         disk_base_url='/'.join([app.config['STUDENT_DISK_URL'], student_id, class_id, task_id, '']),
         class_id=class_id,
         task_id=task_id,
-        computers=computer_list,
-        backing_images=sorted(backing_images),
+        computers=sorted((c for c in computer_list if 'disk_urls' in c), key=lambda c: c['name']),
+        backing_files={fmt: sorted(images) for fmt, images in backing_files.items()},
         lang='sl' if lang == 'si' else lang, # TODO s/si/sl in all tasks (and maybe elsewhere)
         openstack=openstackCreated,
         instructions=jinja2.Template(instructions),
diff --git a/kpov_judge/web/kpov_judge/templates/task_greeting.html b/kpov_judge/web/kpov_judge/templates/task_greeting.html
index 49288e3..ee27853 100644
--- a/kpov_judge/web/kpov_judge/templates/task_greeting.html
+++ b/kpov_judge/web/kpov_judge/templates/task_greeting.html
@@ -56,37 +56,47 @@ section > ul {
 <section class="data">
 <section class="disks">
 <h1>{{ _('Računalniki') }}</h1>
+{% if computers %}
 <p>
-{{ _('Slike diskov za to nalogo:') }}
+{{ _('Za GNS3 uporabite slike v formatu qcow2 ali VMDK, za katere rabite še zaledne datoteke (angl. <em lang="en">backing files</em>). VirtualBox ne podpira zalednih datotek, zato so slike VDI precej večje. Namesto prenosa lahko slike VDI <a href="">dobite iz formata qcow2 ali VMDK</a>.') }}
 
 {% for c in computers %}
 <ul>
   <li>{{c['name']}}
-  {% if 'disk_urls' in c %}
   <ul>
-    {% for u in c['disk_urls'] %}
-    <li><a href="{{disk_base_url+u['file']}}">{{u['name']}}</a>
+    {% for name, disk in c['disk_urls'].items() %}
+    <li>{{name}} [ {% for fmt in disk['formats'] %}<a href="{{disk_base_url+disk[fmt][0]}}">{{fmt}}</a> {% endfor %}]
+    {% else %}
     {% endfor %}
   </ul>
-  {% else %}
-    <br>{{ _('Slike navideznih diskov so v izdelavi in bodo kmalu na voljo.') }}
-  {% endif %}
   </ul>
 {% endfor %}
 
+{% if backing_files %}
 <section>
-{% if backing_images %}
-<h1>{{ _('Osnovne slike') }}</h1>
+<h1>{{ _('Zaledne datoteke') }}</h1>
 <p>
-{{ _('Te slike so enake za vse naloge. Prenesite samo tiste, ki jih še nimate.') }}
+{{ _('Vsako od spodnjih slik prenesete samo pri prvi nalogi, v kateri se pojavi. Vse datoteke morajo biti v istem imeniku.') }}
 
 <ul>
-{% for b in backing_images %}
-  <li><a href="{{disk_base_url+b}}">{{b}}</a></li>
+{% for fmt, images in backing_files.items() %}
+  {% if images %}
+  <li>{{fmt}}
+  <ul>
+    {% for image in images %}
+    <li><a href="{{disk_base_url+image}}">{{image}}</a></li>
+    {% endfor %}
+  </ul>
+  </li>
+  {% endif %}
 {% endfor %}
 </ul>
-{% endif %}
 </section>
+{% endif %}
+
+{% else %}
+  <p>{{ _('Slike navideznih diskov so v izdelavi in bodo kmalu na voljo.') }}
+{% endif %}
 </section>
 
 <section>