#!/usr/bin/env python3

# Formatting of ELF note sections:
# https://docs.oracle.com/cd/E23824_01/html/819-0690/chapter6-18048.html

# The note description field will be packed according to the following struct
#
# struct signature {
#   uint16_t version;
#   uint8_t sig_algo;
#   uint8_t hash_algo;
#   uint32_t length;
#   uint8_t data[];
# }

import re
import sys
import struct
import subprocess

from enum import Enum


if len(sys.argv) != 5:
    print(f"Usage: {sys.argv[0]} <sig-algo> <hash-algo> <livepatch> <signature>", file=sys.stderr)
    sys.exit(1)

(sig_algo, hash_algo, livepatch, signature_binary_file) = sys.argv[1:]

NAME = "Xen"
SECTION_NAME = f".note.{NAME}.signature"
SIGNATURE_FILE = "signature.dat"
ALIGN = 4


class HashAlgorithm(Enum):
    sha256 = 0


class SignatureAlgorithm(Enum):
    rsa = 0


def paddedLength(string):
    return -ALIGN * (-len(string) // ALIGN)

def assertPadded(string, message):
    if len(string) % ALIGN != 0:
        print(message, file=sys.stderr)
        sys.exit(1)

def packV1(sig_algo, hash_algo, signature):
    contents = struct.pack("H", 1)  # Version (uint16_t)
    contents += struct.pack("B", SignatureAlgorithm[sig_algo].value)  # SignatureAlgorithm (uint8_t)
    contents += struct.pack("B", HashAlgorithm[hash_algo].value)  # HashAlgorithm (uint8_t)
    contents += struct.pack("I", len(signature))  # Length (uint32_t)
    contents += struct.pack(f"{len(signature)}s", signature)  # Data (uint8_t[])

    assertPadded(contents, "packV1 contents size is not aligned")
    return contents

def writeNoteFile(filename, name, sig_algo, hash_algo, signature):
    name += b"\x00"
    contents = packV1(sig_algo, hash_algo, signature)
    with open(filename, "wb") as fd:
        fd.write(struct.pack("III", len(name), len(contents), 0))
        fd.write(struct.pack(f"{paddedLength(name)}s", name))
        fd.write(struct.pack(f"{len(contents)}s", contents))


# Validate signing algorithm
if not hasattr(SignatureAlgorithm, sig_algo):
    print(f"Unrecognised signing algorithm: {sig_algo}", file=sys.stderr)
    sys.exit(1)

# Validate hash algorithm
if not hasattr(HashAlgorithm, hash_algo):
    print(f"Unrecognised hash algorithm: {hash_algo}", file=sys.stderr)
    sys.exit(1)

# Read signature from binary file
with open(signature_binary_file, "rb") as sig_file:
    signature = sig_file.read()

print("Removing any existing/conflicting section(s)")
subprocess.call(["objcopy", f"--remove-section={SECTION_NAME}", livepatch])

print("Adding the signature to the livepatch")
writeNoteFile(SIGNATURE_FILE, bytes(NAME, "UTF-8"), sig_algo, hash_algo, signature)
subprocess.check_call(["objcopy", "--add-section", f"{SECTION_NAME}={SIGNATURE_FILE}", livepatch])

print(f"Signature written to .note.{NAME}.signature")
