#!/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 import pyssh import paramiko class SSHGuestFs(): def __init__(hostname, path, username, password): return_results = {} client = paramiko.SSHClient() client.load_system_host_keys() client.set_missing_host_key_policy(paramiko.AutoAddPolicy()) client.connect(hostname, username=username, password=password) self.conn = client self.mountdir = path def chmod(self, mode, path): self.conn.exec_command('chmod {} "{}"'.format(mode, os.path.join(self.mountdir, path))) def chown(self, owner, group, path) self.conn.exec_command('chown {}.{} "{}"'.format(owner, group, os.path.join(self.mountdir, path)) def command(self, arguments): self.conn.exec_command(arguments) def cp(self, src, dest): src = os.path.join(self.mountdir, src) dst = os.path.join(self.mountdir, dst) self.conn.exec_command('cp "{}" "{}"'.format(src, dest)) def cp_a(self, src, dest): src = os.path.join(self.mountdir, src) dst = os.path.join(self.mountdir, dst) self.conn.exec_command('cp -a "{}" "{}"'.format(src, dest)) def cp_r(self, src, dest): src = os.path.join(self.mountdir, src) dst = os.path.join(self.mountdir, dst) self.conn.exec_command('cp -r "{}" "{}"'.format(src, dest)) def dd(self, src, dest): src = os.path.join(self.mountdir, src) dst = os.path.join(self.mountdir, dst) self.conn.exec_command('dd if="{}" of="{}"'.format(src, dest)) def df(self): stdin, stdout, stderr = self.conn.exec_command('df') return stdin.read() def download(self, remotefilename, filename): path = os.path.join(self.mountdir, remotefilename) stdin, stdout, stderr = self.conn.exec_command('dd if="{}"'.format(path)) with open(filename, 'w') as f: while data = stdin.read(4096): f.write(data) def download_offset (self, remotefilename, filename, offset, size): path = os.path.join(self.mountdir, remotefilename) stdin, stdout, stderr = self.conn.exec_command('dd bs=1 skip={} count={} if="{}"'.format( offset, size ,path)) with open(filename, 'w') as f: while data = stdin.read(4096): f.write(data) def du(self, path) path = os.path.join(self.mountdir, path) stdin, stdout, stderr = self.conn.exec_command('du "{}"'.format(path)) return stdin.read() def equal(self, file1, file2) src = os.path.join(self.mountdir, src) dst = os.path.join(self.mountdir, dst) pass def file(self, path) path = os.path.join(self.mountdir, path) stdin, stdout, stderr = self.conn.exec_command('file "{}"'.format(path)) return stdin.read() def getxattrs(self, path) pass #path = os.path.join(self.mountdir, path) #stdin, stdout, stderr = self.conn.exec_command('du "{}"'.format(path)) #return stdin.read() def mv (self, src, dest): src = os.path.join(self.mountdir, src) dst = os.path.join(self.mountdir, dst) self.conn.exec_command('mv "{}" "{}"'.format(src, dest)) def mkdir (self, path): path = os.path.join(self.mountdir, path) sftp_client = self.conn.open_sftp() sftp_client.mkdir(path) def read_file (self, path): path = os.path.join(self.mountdir, path) sftp_client = self.conn.open_sftp() f = sftp_client.file(path, mode='r', bufsize=-1) s = f.read() f.close() return s def readdir (self, dir): path = os.path.join(self.mountdir, path) sftp_client = self.conn.open_sftp() return sftp_client.listdir(path) def readlink (self, path): path = os.path.join(self.mountdir, path) stdin, stdout, stderr = self.conn.exec_command('readlink "{}"'.format(path)) return stdin.read() def rename (self, oldpath, newpath): return self.mv(oldpath, newpath) def rm (self, path): path = os.path.join(self.mountdir, path) stdin, stdout, stderr = self.conn.exec_command('rm "{}"'.format(path)) def rm_rf (self, path): path = os.path.join(self.mountdir, path) stdin, stdout, stderr = self.conn.exec_command('rm -rf "{}"'.format(path)) def rmdir (self, path): path = os.path.join(self.mountdir, path) stdin, stdout, stderr = self.conn.exec_command('rmdir "{}"'.format(path)) def setxattr (self, xattr, val, vallen, path): pass def write (self, path, content): """This call creates a file called "path". The content of the file is the string "content" (which can contain any 8 bit data). See also "g.write_append". """ path = os.path.join(self.mountdir, path) sftp_client = self.conn.open_sftp() f = sftp_client.file(path, mode='w', bufsize=-1) f.write(content) f.close() return r def write_append (self, path, content): """This call appends "content" to the end of file "path". If "path" does not exist, then a new file is created. See also "g.write". """ path = os.path.join(self.mountdir, path) sftp_client = self.conn.open_sftp() f = sftp_client.file(path, mode='a', bufsize=-1) f.write(content) f.close() 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()