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

import ipaddress
import os
import socketserver
import json
import re
import socket
import sys

PORT = int(os.environ.get("GSTFSD_PORT", 16510))
ADDRESS = os.environ.get("GSTFSD_BIND_HOST", "127.0.0.1")


def is_loopback(host):
    try:
        return ipaddress.ip_address(host).is_loopback
    except ValueError:
        return host == "localhost"


# gstfsd is an unauthenticated root service that rewrites any VM's
# /etc/shadow and authorized_keys, so it only listens on loopback unless
# explicitly allowed (ROADMAP S-08).
if not is_loopback(ADDRESS) and os.environ.get("GSTFSD_ALLOW_REMOTE") != "1":
    sys.stderr.write(
        "gstfsd: refusing to listen on %s (unauthenticated root service). "
        "Set GSTFSD_ALLOW_REMOTE=1 to allow it on a trusted network.\n" % ADDRESS
    )
    sys.exit(1)

import guestfs  # noqa: E402


class MyTCPServer(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    address_family = socket.AF_INET6 if ":" in ADDRESS else socket.AF_INET


class MyTCPServerHandler(socketserver.StreamRequestHandler):
    timeout = 30  # seconds; an idle client must not hold a thread forever

    def handle(self):
        # recive data
        d = self.request.recv(1024).strip()
        data = json.loads(d)

        # GuestFS
        gfs = guestfs.GuestFS(python_return_dict=True)
        try:
            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(0o640, '/etc/shadow')
                            self.request.sendall(json.dumps({'return': 'success'}).encode())
                        if data['action'] == 'publickey':
                            if not gfs.is_dir('/root/.ssh'):
                                gfs.mkdir('/root/.ssh')
                                gfs.chmod(0o700, "/root/.ssh")
                            # Add the key; keep the ones already there.
                            keys_path = '/root/.ssh/authorized_keys'
                            keys = gfs.cat(keys_path) if gfs.is_file(keys_path) else ''
                            new_key = data['key'].strip()
                            if new_key not in keys.splitlines():
                                if keys and not keys.endswith('\n'):
                                    keys += '\n'
                                gfs.write(keys_path, keys + new_key + '\n')
                            gfs.chmod(0o600, '/root/.ssh/authorized_keys')
                            self.request.sendall(json.dumps({'return': 'success'}).encode())
                    gfs.umount(part)
                except RuntimeError:
                    pass
            gfs.shutdown()
            gfs.close()
        except Exception as err:
            self.request.sendall(bytes(json.dumps({'return': 'error', 'message': str(err)}).encode()))


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