Initial commit

git-svn-id: https://svn.lusy.fri.uni-lj.si/kpov-public-svn/kpov-public@1 5cf9fbd1-b2bc-434c-b4b7-e852f4f63414

1 files changed, 164 insertions, 0 deletions

diff --git a/kpov_judge/create_disk_images.py b/kpov_judge/create_disk_images.py
new file mode 100755
index 0000000..168aa5a
--- /dev/null
+++ b/kpov_judge/create_disk_images.py
@@ -0,0 +1,164 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import pymongo
+import sys
+import inspect
+import kpov_random_helpers
+import settings
+import guestfs
+import os
+import glob
+import subprocess
+import fcntl
+
+def get_prepare_disks(db, task_id):
+    prepare_disks_source = db.prepare_disks.find_one({'task_id':task_id})['source']
+    prepare_disks_code = compile(prepare_disks_source, 'prepare_disks.py', 'exec')
+    exec(prepare_disks_code)
+    return prepare_disks
+
+def create_snapshot(task_id, student_id, disk_name, overwrite = True, cow = False):
+    print(os.path.join(settings.DISK_TEMPLATE_PATH, task_id, disk_name) + '.*')
+    template_paths = glob.glob(os.path.join(settings.DISK_TEMPLATE_PATH, task_id, disk_name) + '.*')
+    filtered_paths = filter((lambda x: os.path.splitext(x)[1] == '.' + settings.STUDENT_DISK_FORMAT), template_paths)
+    if filtered_paths:
+      template_path = filtered_paths[0]
+    else:
+      template_path = template_paths[0]
+    if cow:
+        d = os.path.join(student_id, task_id, disk_name) + os.path.splitext(template_path)[1]
+    else:
+        d = os.path.join(student_id, task_id, disk_name) + '.qcow2'
+    try:
+        os.makedirs(os.path.join(settings.STUDENT_DISK_PATH, student_id, task_id))
+    except:
+        pass
+    disk_file = os.path.join(settings.STUDENT_DISK_PATH, d)
+    if overwrite or not os.path.exists(disk_file):
+        if cow:
+            subprocess.call(['cp', '--reflink=always', template_path, disk_file])
+        else:
+            subprocess.call(['qemu-img', 'create', '-f', 'qcow2', '-o', 'backing_file=' + template_path,
+                     disk_file])
+    return d
+
+def publish_snapshot(d):
+    if os.path.splitext(d)[1][1:] != settings.STUDENT_DISK_FORMAT:
+        snap_name = os.path.splitext(d)[0] + '.' + settings.STUDENT_DISK_FORMAT
+        disk_file = os.path.join(settings.STUDENT_DISK_PATH, d)
+        snap_file = os.path.join(settings.STUDENT_DISK_PATH, snap_name)
+        subprocess.call(['qemu-img', 'convert', '-f', 'qcow2', '-O', settings.STUDENT_DISK_FORMAT, disk_file,
+            snap_file])
+        url = settings.STUDENT_DISK_URL + snap_name
+    else:
+        url = settings.STUDENT_DISK_URL + d
+    return url
+  
+if __name__ == '__main__':
+    if len(sys.argv) != 1:
+        print "Usage: {0}"
+        print "Create the neccessarry disk images"
+    db = pymongo.Connection(settings.DB_HOST).kpov
+    try:
+        db.authenticate(settings.USERNAME, settings.PASSWORD)
+    except Exception, e:
+        print str(e)
+    dev_prefix = settings.GUESTFS_DEV_PREFIX
+    l = db.student_computers.find({"disk_urls": {"$exists": False}})
+    computers_by_task_student = dict()
+    for computer in l:
+        student_id, task_id = computer['student_id'], computer['task_id']
+        if (task_id, student_id) not in computers_by_task_student:
+            computers_by_task_student[(task_id, student_id)] = list()
+        computers_by_task_student[(task_id, student_id)].append(computer)
+    for (task_id, student_id), computers in computers_by_task_student.iteritems():
+    	l = db.student_computers.find_one({'task_id': task_id, 'student_id':student_id, "disk_urls": {"$exists": False}})
+        if l is None:
+            continue
+        lock_file = os.path.join(settings.STUDENT_LOCKFILE_PATH, 
+            '{0}-{1}.lock'.format(student_id, task_id))
+        lock_fp = open(lock_file, 'w')
+        try:
+            fcntl.lockf(lock_fp, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        except IOError:
+            continue
+        params = db.task_params.find_one({'task_id': task_id, 'student_id': student_id})['params']
+        prepare_disks = get_prepare_disks(db, task_id)
+        # tule odpri, ustvari snapshote za vajo
+        templates = dict()
+        all_disks = dict()
+        parts = dict()
+        for computer in computers:
+            lock_fp.write('creating computer ' + computer['name'] + '\n')
+            all_disks[computer['name']] = dict()
+            manual_disks = list()
+            this_computers_disks = set()
+            try_automount = False
+            g = guestfs.GuestFS()
+            for disk in computer['disks']:
+                lock_fp.write("register " + disk['name'] + '\n')
+                snap = create_snapshot(task_id, student_id, disk['name'],
+                    cow = settings.STUDENT_DISK_COW)
+                snap_file = os.path.join(settings.STUDENT_DISK_PATH, 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']:
+                        manual_disks.append((dev_prefix + p['dev'],
+                            p['path'], p.get('options', None)))
+                else:
+                    try_automount = True
+                templates[disk['name']] = g
+                all_disks[computer['name']][disk['name']] = snap
+            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")
+        prepare_disks(templates, 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 comp_name, d_dict in all_disks.iteritems():
+            disk_urls = list()
+            for d_name, d in d_dict.iteritems():
+                lock_fp.write('publishing '+ str(d) + '\n')
+                url = publish_snapshot(d)
+                lock_fp.write('published as '+ url + '\n')
+                disk_urls.append({'name': d_name, 'url': url})
+            lock_fp.write('urls: '+ str(disk_urls) + '\n')
+            l = db.student_computers.update({
+                "disk_urls": {"$exists": False}, 
+                'student_id': student_id, 
+                'task_id': task_id,
+                'name': comp_name},
+                {'$set': { 'disk_urls': disk_urls }})
+        # print "done for ", student_id, task_id
+        os.unlink(lock_file)
+        lock_fp.close()