#!/usr/bin/env python
#
# gstfsd - WebVirtCloud daemon for managing VM's filesystem
#

import SocketServer
import json
import guestfs
import re
import os
from jwcrypto import jws, jwk, jwe

PORT = 16510
ADDRESS = "0.0.0.0"
SECRET = None


class MyTCPServer(SocketServer.ThreadingTCPServer):
    allow_reuse_address = True


class MyTCPServerHandler(SocketServer.BaseRequestHandler):
    def handle(self):
        # recive data and check authentcation
        try:
            signed_data = jws.JWS()
            signed_data.deserialize(self.request.recv(4096).strip())
            signed_data.verify(SECRET, "HS512")

            encrypted_data = jwe.JWE(algs=["A256KW", "A256CBC-HS512"])
            encrypted_data.deserialize(signed_data.payload)
            encrypted_data.decrypt(SECRET)

            data = json.loads(encrypted_data.plaintext)

            # GuestFS
            gfs = guestfs.GuestFS(python_return_dict=True)
            gfs.add_domain(data['vname'])
            gfs.launch()
            parts = gfs.list_partitions()
            for part in parts:
                try:
                    gfs.mount(part, '/')
                    if gfs.is_file('/etc/shadow'):
                        if data['action'] == 'password':
                            file_shadow = gfs.cat('/etc/shadow')
                            new_root_hash = "root:" + data['passwd'] + ":"
                            file_shadow_new = re.sub('^root:.*?:', new_root_hash, file_shadow)
                            gfs.write('/etc/shadow', file_shadow_new)
                            gfs.chmod(640, '/etc/shadow')
                            self.request.sendall(json.dumps({'return': 'success'}))
                        if data['action'] == 'publickey':
                            if not gfs.is_dir('/root/.ssh'):
                                gfs.mkdir('/root/.ssh')
                                gfs.chmod(0700, "/root/.ssh")
                            gfs.write('/root/.ssh/authorized_keys', data['key'])
                            gfs.chmod(0600, '/root/.ssh/authorized_keys')
                            self.request.sendall(json.dumps({'return': 'success'}))
                    gfs.umount(part)
                except RuntimeError:
                    pass
            gfs.shutdown()
            gfs.close()
        # we check signature before trying to decrypt so jwe.InvalidJWEData should not be raised ever
        except (jws.InvalidJWSObject, jwe.InvalidJWEData, RuntimeError, ValueError) as err:
            self.request.sendall(json.dumps({'return': 'error', 'message': err.message}))

        except jws.InvalidJWSSignature as err:
            self.request.sendall(json.dumps({'return': 'error', 'message': (
                "Fail to verify request signature. Check if you have imported "
                "the key (/var/lib/gstfsd/SECRET) in WebVirtCloud"
            )}))

if not os.path.isfile("/var/lib/gstfsd/SECRET"):
    try:
        os.mkdir("/var/lib/gstfsd")
    except OSError as error:
        if error.errno != 17: # File exists
            raise
    os.chmod("/var/lib/gstfsd", 0700)
    with open("/var/lib/gstfsd/SECRET", 'w') as f:
        f.write(jwk.JWK(generate='oct', size=256).export())
    os.chmod("/var/lib/gstfsd/SECRET", 0600)

with open("/var/lib/gstfsd/SECRET") as f:
    SECRET = jwk.JWK(**json.load(f))

server = MyTCPServer((ADDRESS, PORT), MyTCPServerHandler)
server.serve_forever()