#!/usr/bin/env python3
# Based on:
#  * https://github.com/newsoft/dbx-keygen-linux
#  * https://github.com/newsoft/dbx-keygen-windows

import binascii
import hashlib
import hmac
import json
import os
import struct
import sys
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import PBKDF2
from pprint import pprint

DIGESTMOD = hashlib.md5

def _trace(msg, *args):
    if os.environ.get("DEBUG"):
        if args:
            msg = msg % args
        print("DEBUG:", msg, file=sys.stderr)

def b64encode(buf):
    return binascii.b2a_base64(buf).decode().replace("\n", "")

def b64decode(buf):
    return binascii.a2b_base64(buf)

class KeyStore(object):
    CLIENT_KEY_NAME = "Client"

    def __init__(self):
        self.parsers = {0: ParserVersion0()}
        self.user_hmac_keys = {k: v.USER_HMAC_KEY for k, v in self.parsers.items()}

    def get_user_key(self):
        return self.get_keyv(self.CLIENT_KEY_NAME)[1]

    def get_database_key(self):
        version, user_key = self.get_keyv(self.CLIENT_KEY_NAME)
        return self.parsers[version].derive_database_key(user_key)

    def backup(self):
        return {k: self.get_keyv(k)[1] for (k, v) in self.enum_raw()}

    def restore(self, data):
        for k, v in data.items():
            self.put_keyv(k, v)

class FileKeyStore(KeyStore):
    FILE_HMAC_KEY = b'\x8f\xf4\xf2\xbb\xad\xe9\x47\xea\x1f\xdf\x69\x6d\x80\x5b\x35\x3e'

    @classmethod
    def list_instances(self):
        return [name
                for name in os.listdir(os.path.expanduser("~/.dropbox"))
                if name.startswith("instance")]

    def __init__(self, instance_name):
        super().__init__()
        self.appdata_path = os.path.join(os.path.expanduser("~/.dropbox"), instance_name)
        self.keyfile_path = os.path.join(self.appdata_path, "hostkeys")
        self.file_hmac_keys = {0: self.FILE_HMAC_KEY}
        self.file_key = self._uid_to_key(self.unique_id(self.appdata_path))
        self.file_iv = b'\x6c\x07\x38\x01\x34\x24\x73\x58\x03\xff\x72\x69\x33\x13\x61\x51'
        self._vers = 0
        self._dict = {}
        self._loaded = False

    @property
    def instance_name(self):
        return os.path.basename(self.appdata_path)

    def unique_id(self, appdata_path):
        raise NotImplementedError()

    def _uid_to_key(self, uid):
        magic = 'ia9%sXa|ui20' % uid
        return hashlib.md5(magic.encode()).digest()

    def protect(self, data):
        encryptor = AES.new(key=self.file_key,
                            mode=AES.MODE_CBC,
                            IV=self.file_iv)
        bs = encryptor.block_size
        padn = bs - (len(data) % bs)
        return encryptor.encrypt(data + bytes([padn]*padn))

    def unprotect(self, data):
        decryptor = AES.new(key=self.file_key,
                            mode=AES.MODE_CBC,
                            IV=self.file_iv)
        data = decryptor.decrypt(data)
        padn = data[-1]
        return data[:-padn]

    def unpack_verify(self, packed_buf, hmac_keys, protected=False):
        header_fmt = "=B"
        header_len = struct.calcsize(header_fmt)

        (version,) = struct.unpack_from(header_fmt, packed_buf)
        if version not in hmac_keys:
            raise Exception("Unknown payload version %r" % version)

        hm = hmac.new(hmac_keys[version], digestmod=DIGESTMOD)
        ds = hm.digest_size
        if len(packed_buf) <= ds:
            raise Exception("Buffer shorter than digest size")

        hm.update(packed_buf[:-ds])
        if hm.digest() != packed_buf[-ds:]:
            raise Exception('Bad digest')

        payload = packed_buf[header_len:-ds]
        if protected:
            payload = self.unprotect(payload)

        return (version, payload)

    def pack_sign(self, version, payload, hmac_keys, protected=False):
        if version not in hmac_keys:
            raise Exception("Unknown payload version %r" % version)

        if protected:
            payload = self.protect(payload)
        header = struct.pack("=B", version)
        hm = hmac.new(hmac_keys[version], digestmod=DIGESTMOD)
        hm.update(header)
        hm.update(payload)
        return header + payload + hm.digest()

    def _read_keyfile(self, path):
        packed_buf = open(path, "rb").read()
        _trace("packed_buf: %r", packed_buf)
        return self.unpack_verify(packed_buf, self.file_hmac_keys,
                                              protected=True)

    def _write_keyfile(self, path, version, payload):
        packed_buf = self.pack_sign(version, payload, self.file_hmac_keys,
                                    protected=True)
        _trace("packed_buf: %r", packed_buf)
        open(path, "wb").write(packed_buf)

    def _load(self):
        if not self._loaded:
            (version, payload) = self._read_keyfile(self.keyfile_path)
            self._vers = version
            if hasattr(payload, "decode"):
                payload = payload.decode()
            self._dict = json.loads(payload)
            self._loaded = True

    def _store(self):
        self._write_keyfile(self.keyfile_path,
                            version=self._vers,
                            payload=json.dumps(self._dict).encode())

    def enum_raw(self):
        self._load()
        return self._dict.items()

    def get_raw(self, name):
        self._load()
        return b64decode(self._dict[name])

    def put_raw(self, name, value):
        self._dict[name] = b64encode(value)
        self._store()

    def get_keyv(self, name):
        return self.unpack_verify(self.get_raw(name), self.user_hmac_keys)

    def put_keyv(self, name, value, version=0):
        self.put_raw(name, self.pack_sign(version, value, self.user_hmac_keys))

class LinuxKeyStore(FileKeyStore):
    @staticmethod
    def linux_statvfs(path):
        import ctypes
        import ctypes.util

        class struct_statvfs(ctypes.Structure):
            fsblkcnt_t = ctypes.c_ulong
            fsfilcnt_t = ctypes.c_ulong
            _fields_ = [
                ('f_bsize',     ctypes.c_ulong),
                ('f_frsize',    ctypes.c_ulong),
                ('f_blocks',    fsblkcnt_t),
                ('f_bfree',     fsblkcnt_t),
                ('f_bavail',    fsblkcnt_t),
                ('f_files',     fsfilcnt_t),
                ('f_ffree',     fsfilcnt_t),
                ('f_favail',    fsfilcnt_t),
                ('f_fsid',      ctypes.c_ulong),
                ('f_flag',      ctypes.c_ulong),
                ('f_namemax',   ctypes.c_ulong),
                ('_padding_',   ctypes.c_int * 128),
            ]

        libc = ctypes.cdll[ctypes.util.find_library("c")]
        _statvfs = libc.statvfs
        _statvfs.restype = ctypes.c_int
        _statvfs.argtypes = [ctypes.c_char_p, ctypes.POINTER(struct_statvfs)]
        buf = struct_statvfs()
        ret = _statvfs(path.encode(), ctypes.byref(buf))
        if ret < 0:
            raise Exception('statvfs failed with retval %s' % (ret,))
        return buf

    def unique_id(self, path):
        inode = os.stat(path).st_ino
        fsid = self.linux_statvfs(path).f_fsid
        _trace("unique ID for %s: {inode=%d, fsid=%d}", path, inode, fsid)
        return '%d_%d' % (inode, fsid)

class WindowsKeyStore(KeyStore):
    CRYPTPROTECT_UI_FORBIDDEN = 0x01
    CRYPTPROTECT_LOCAL_MACHINE = 0x04

    def __init__(self, instance_name):
        super().__init__()
        if instance_name == "instance_db":
            instance_name = "ks"
        elif instance_name.startswith("instance"):
            instance_name = "ks" + instance_name[8:]
        self.instance_name = instance_name
        self.reg_base = "SOFTWARE\\Dropbox"
        self.reg_path = "%s\\%s" % (self.reg_base, instance_name)

    def protect(self, data_in, extra_entropy=None):
        import win32crypt
        desc = ""
        flags = self.CRYPTPROTECT_UI_FORBIDDEN
        data_out = win32crypt.CryptProtectData(data_in, desc, extra_entropy, None, None, flags)
        return data_in

    def unprotect(self, data_in, extra_entropy=None):
        import win32crypt
        flags = self.CRYPTPROTECT_UI_FORBIDDEN
        (desc, data_out) = win32crypt.CryptUnprotectData(data_in, extra_entropy, None, None, flags)
        _trace("description: %r", desc)
        return data_out

    def unpack_verify(self, packed_buf, hmac_keys, protected=True):
        header_fmt = "<LL"
        header_len = struct.calcsize(header_fmt)

        (version, payload_len) = struct.unpack_from(header_fmt, packed_buf)
        if version not in hmac_keys:
            raise Exception("Unknown payload version %r" % version)

        hm = hmac.new(hmac_keys[version], digestmod=DIGESTMOD)
        ds = hm.digest_size
        if len(packed_buf) <= ds:
            raise Exception("Buffer shorter than digest size")
        if len(packed_buf) != header_len + payload_len + ds:
            raise Exception("Buffer length does not match payload_len")

        hm.update(packed_buf[:-ds])
        if hm.digest() != packed_buf[-ds:]:
            raise Exception("HMAC mismatch")

        payload = packed_buf[header_len:-ds]
        if protected:
            payload = self.unprotect(payload, hmac_keys[version])

        return (version, payload)

    def pack_sign(self, version, payload, hmac_keys, protected=True):
        if version not in hmac_keys:
            raise Exception("Unknown payload version %r" % version)

        if protected:
            payload = self.protect(payload, hmac_keys[version])
        header = struct.pack("<LL", version, len(payload))
        hm = hmac.new(hmac_keys[version], digestmod=DIGESTMOD)
        hm.update(header)
        hm.update(payload)
        return header + payload + hm.digest()

    @classmethod
    def list_instances(self):
        import winreg
        hkey = winreg.OpenKey(winreg.HKEY_CURRENT_USER, self.reg_base)
        i = 0
        while True:
            try:
                subkey = winreg.EnumKey(hkey, i)
                if subkey.startswith("ks"):
                    yield subkey
                i += 1
            except winreg.error:
                break

    def enum_raw(self):
        import winreg
        hkey = winreg.OpenKey(winreg.HKEY_CURRENT_USER, self.reg_path)
        i = 0
        while True:
            try:
                (name, data, type) = winreg.EnumValue(hkey, i)
                yield (name, data)
                i += 1
            except winreg.error:
                break

    def get_raw(self, name):
        import winreg
        hkey = winreg.OpenKey(winreg.HKEY_CURRENT_USER, self.reg_path)
        (data, type) = winreg.QueryValueEx(hkey, name)
        return data[:-1] # remove null byte

    def put_raw(self, name, value):
        import winreg
        hkey = winreg.OpenKey(winreg.HKEY_CURRENT_USER, self.reg_path)
        winreg.SetValueEx(hkey, name, None, winreg.REG_BINARY, value + b"\0")

    def get_keyv(self, name):
        return self.unpack_verify(self.get_raw(name), self.user_hmac_keys,
                                  protected=(not name.endswith("-p")))

    def put_keyv(self, name, value, version=0):
        self.put_raw(name, self.pack_sign(version, value, self.user_hmac_keys,
                                          protected=(not name.endswith("-p"))))

# ---------------------------------------------

class ParserVersion0(object):
    USER_HMAC_KEY = b'\xd1\x14\xa5\x52\x12\x65\x5f\x74\xbd\x77\x2e\x37\xe6\x4a\xee\x9b'
    APP_KEY = b'\x0d\x63\x8c\x09\x2e\x8b\x82\xfc\x45\x28\x83\xf9\x5f\x35\x5b\x8e'
    APP_IV = b'\xd8\x9b\x43\x1f\xb6\x1d\xde\x1a\xfd\xa4\xb7\xf9\xf4\xb8\x0d\x05'
    APP_KDF_ITER = 1066
    USER_KEYLEN = 16
    DB_KEYLEN = 16

    def derive_database_key(self, user_key):
        return PBKDF2(password=user_key,
                      salt=self.APP_KEY,
                      dkLen=self.DB_KEYLEN,
                      count=self.APP_KDF_ITER)

# ---------------------------------------------

def get_keystore():
    name = os.uname().sysname
    if name == "Linux":
        return LinuxKeyStore
    else:
        raise Exception("Unsupported OS %r" % name)

def iter_keystores():
    klass = get_keystore()
    for path in klass.list_instances():
        yield klass(path)

def do_help(argv):
    print("Commands:")
    print()
    print("  show")
    print("  dump > keys.json")
    print("  load < keys.json")

def do_show(argv):
    for ks in iter_keystores():
        print("Instance:", ks.instance_name)
        print(" * User key:", b64encode(ks.get_user_key()))
        print(" * Database key:", b64encode(ks.get_database_key()))

def do_dump(argv):
    d = {}
    for ks in iter_keystores():
        di = ks.backup()
        di = {k: b64encode(v) for k, v in di.items()}
        d[ks.instance_name] = di
    d = json.dumps(d)
    f = open(argv[0], "w") if argv else sys.stdout
    print(d, file=f)

def do_load(argv):
    f = open(argv[0], "r") if argv else sys.stdin
    d = f.read()
    d = json.loads(d)
    for name, di in d.items():
        di = {k: b64decode(v) for k, v in di.items()}
        ks = get_keystore()(name)
        ks.restore(di)
        print("restored to %r", ks.instance_name)

try:
    cmd, *args = sys.argv[1:]
except ValueError:
    cmd, args = "help", []

try:
    func = globals()["do_%s" % cmd.replace("-", "_")]
except KeyError:
    exit("unknown command %r", cmd)

func(args)
