#!/usr/bin/env python3

import argparse
import functools
import hmac
import json
import plistlib
import subprocess


memoize_forever = functools.lru_cache(maxsize=None)


# Analogous to GetDeterministicMachineSpecificId from device_id_mac.cc,
# except we shell out to `ioreg(1)` instead of using the C API
@memoize_forever
def get_machine_id():
    xml = subprocess.check_output('ioreg -c IOPlatformExpertDevice -d 1 -r -a'.split())
    data = plistlib.loads(xml)
    return data[0]['IOPlatformUUID']


# Chromium uses the empty string for a seed. Chrome proper uses a blob stored
# in the resource bundle with key IDR_PREF_HASH_SEED_BIN which is loaded from
# the non-public file resources\settings_internal\pref_hash_seed.bin. This is
# quite stupid, because the seed can trivially be extracted back out of the
# resource bundle. Note that, as of 2019-08, there is no support for rolling
# the seed, and this value appears to have been used consistently across
# various versions and platforms.
SEED = bytes.fromhex(
    'e748f336d85ea5f9dcdf25d8f347a65b4cdf667600f02df6724a2af18a212d26'
    'b788a25086910cf3a90313696871f3dc05823730c91df8ba5c4fd9c884b505a8'
)


# port of GetDigestString from pref_hash_calculator.cc
def get_digest_string(key: bytes, message: bytes) -> str:
    dgst = hmac.new(key=key, msg=message, digestmod='sha256')
    return dgst.hexdigest().upper()


# loose port of CopyWithoutEmptyChildren and friends from values.cc
def copy_without_empty_children(value):
    if isinstance(value, list):
        copy = []
        for child in value:
            child_copy = copy_without_empty_children(child)
            if child_copy is not None:
                copy.append(child_copy)
        return copy if copy else None
    elif isinstance(value, dict):
        copy = {}
        for k, v in value.items():
            v_copy = copy_without_empty_children(v)
            if v_copy is not None:
                copy[k] = v_copy
        return copy if copy else None
    elif isinstance(value, (type(None), bool, int, float, str, bytes)):
        return value
    else:
        raise TypeError("Can't copy_without_empty_children odd value of type %r" % (type(value),))


# port of JSONWriter (yes, Chrome rolled their own).
# Also JSONStringValueSerializer, which is a thin wrapper around it
def chrome_json_ser(value) -> bytes:
    if value is None:
        return 'null'.encode()
    elif isinstance(value, bool):
        return (('true' if value else 'false').encode())
    elif isinstance(value, int):
        return str(value).encode()
    elif isinstance(value, float):
        real = str(value) # XXX assumes python is same as chrome's NumberToString
        if '.' not in real and 'e' not in real and 'E' not in real:
            real = real + '.0'
        elif real[0] == '.':
            real = '0' + real
        elif real[0] == '-' and real[0] == '.':
            real = '-0' + real[1:]
        return real.encode()
    elif isinstance(value, str):
        # XXX doesn't handle cases where str has illegal unicode
        out = b'"'
        for cp in value:
            special = {
                    '\b': '\\b',
                    '\f': '\\f',
                    '\n': '\\n',
                    '\r': '\\r',
                    '\t': '\\t',
                    '\\': '\\\\',
                    '"': '\\"',
                    '<': '\\u003C',
                    '\u2028': '\\u2028',
                    '\u2029': '\\u2029',
            }.get(cp, None)
            if special is not None:
                out += special.encode()
            elif ord(cp) < 32:
                out += ('\\u00%02x' % (ord(cp),)).encode()
            else:
                out += cp.encode()
        out += b'"'
        return out
    elif isinstance(value, list):
        out = b'['
        first = True
        for v in value:
            if not first:
                out += b','
            out += chrome_json_ser(v)
            first = False
        out += b']'
        return out
    elif isinstance(value, dict):
        out = b'{'
        first = True
        for k, v in value.items():
            if not first:
                out += b','
            out += chrome_json_ser(k)
            out += b':'
            out += chrome_json_ser(v)
            first = False
        out += b'}'
        return out
    elif isinstance(value, bytes):
        raise TypeError("Chrome's JSON library doesn't support bytes")
    else:
        raise TypeError("Can't chrome_json_ser odd value of type %r" % (type(value),))


# port of ValueAsString pref_hash_calculator.cc
def value_as_string(value) -> bytes:
    if isinstance(value, dict):
        value = copy_without_empty_children(value) or {}
    return chrome_json_ser(value)


# port of GetMessage from pref_hash_calculator.cc
def get_message(device_id: str, path: str, value_as_bytes: bytes) -> bytes:
    return device_id.encode() + path.encode() + value_as_bytes


# port of PrefHashCalculator::Calculate from pref_hash_calculator.cc
def calculate(path: str, value) -> str:
    return get_digest_string(
            key=SEED,
            message=get_message(
                device_id=get_machine_id(),
                path=path,
                value_as_bytes=value_as_string(value)))


# Traverse the in-memory representation of a 'Secure Preferences' file and
# replace all the MACs with new ones computed for the current computer.
def recompute_protection(secure_prefs_data):
    # Chrome's prefs system appears to have coded-in handling for whether
    # certain keys are 'ATOMIC' or 'SPLIT'. Duplicating the list here seems
    # clowny, so I don't see a great way of generating the MACs based solely on
    # the corresponding values. However, traversing the already-existing MACs
    # and replacing them seems viable; if we assume that Chrome is the only
    # thing which added keys, then it should have written MACs out in the right
    # way when it added them. In terms of the implementation, I'm guessing a
    # bit, but this seems to handle most (if not all) cases and there's a limit
    # to how much time I want to spend reading Chrome's prefs code.

    def recompute_mac_subtree(parent_keys, mac_subtree, prefs_subtree):
        for k in sorted(mac_subtree.keys()):
            if k not in prefs_subtree:
                # AFAICT Chrome just leaves macs lying around even when the
                # pref is gone. I think the right thing to do is skip them
                # here-- we have no value to recompute a new mac and deleting
                # them seems unnecessarily risky
                continue
            if isinstance(mac_subtree[k], dict):
                recompute_mac_subtree(parent_keys + [k], mac_subtree[k], prefs_subtree[k])
            elif isinstance(mac_subtree[k], str):
                mac_subtree[k] = calculate('.'.join(parent_keys + [k]), prefs_subtree[k])
            else:
                raise TypeError('Weird thing found in protection.macs dict')

    if 'protection' in secure_prefs_data:
        protection = secure_prefs_data['protection']
        if 'macs' in protection:
            recompute_mac_subtree([], protection['macs'], secure_prefs_data)
            if 'super_mac' in protection:
                protection['super_mac'] = calculate('', protection['macs'])


def resign(inpath, outpath):
    with open(inpath, 'r', encoding='utf8') as fp:
        data = json.load(fp)
    recompute_protection(data)
    with open(outpath, 'w', encoding='utf8') as fp:
        json.dump(data, fp)


def main():
    parser = argparse.ArgumentParser(
        description="Resign Chrome's 'Secure Preferences' for the current machine",
    )
    parser.add_argument(
        '--in',
        metavar='FILE_PATH',
        help='path to old Secure Preferences file',
        required=True,
    )
    parser.add_argument(
        '--out',
        metavar='FILE_PATH',
        help='path to write out new Secure Preferences file',
        required=True
    )
    args = parser.parse_args()
    resign(vars(args)['in'], args.out)


if __name__ == '__main__':
    main()