graphene/Pal/src/host/Linux-SGX/signer/pal-sgx-get-token

#!/usr/bin/env python3

import argparse
import os
import socket
import struct
import sys

import aesm_pb2

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from generated_offsets import *


def read_sigstruct(sig):
    """Reading Sigstruct."""
    # field format: (offset, type, value)
    # SGX_ARCH_SIGSTRUCT_
    fields = {
        'date': (SGX_ARCH_SIGSTRUCT_DATE, "<HBB", 'year', 'month', 'day'),
        'modulus': (SGX_ARCH_SIGSTRUCT_MODULUS, "384s", 'modulus'),
        'exponent': (SGX_ARCH_SIGSTRUCT_EXPONENT, "<L", 'exponent'),
        'signature': (SGX_ARCH_SIGSTRUCT_SIGNATURE, "384s", 'signature'),

        'miscs': (SGX_ARCH_SIGSTRUCT_MISCSELECT, "4s", 'miscs'),
        'miscmask': (SGX_ARCH_SIGSTRUCT_MISCSELECT_MASK, "4s", 'miscmask'),
        'attrs': (SGX_ARCH_SIGSTRUCT_ATTRIBUTES, "8s8s", 'flags', 'xfrms'),
        'attrmask': (SGX_ARCH_SIGSTRUCT_ATTRIBUTES_MASK,
                     "8s8s", 'flagmask', 'xfrmmask'),
        'mrenclave': (SGX_ARCH_SIGSTRUCT_ENCLAVE_HASH, "32s", 'mrenclave'),
        'isvprodid': (SGX_ARCH_SIGSTRUCT_ISVPRODID, "<H", 'isvprodid'),
        'isvsvn': (SGX_ARCH_SIGSTRUCT_ISVSVN, "<H", 'isvsvn'),
    }

    attr = dict()
    for field in fields.values():
        values = struct.unpack_from(field[1], sig, field[0])

        for i in range(len(values)):
            attr[field[i + 2]] = values[i]

    return attr


def connect_aesmd(attr):
    """Connect with AESMD."""

    req_msg = aesm_pb2.GetTokenReq()
    req_msg.req.signature = attr['mrenclave']
    req_msg.req.key = attr['modulus']
    req_msg.req.attributes = attr['flags'] + attr['xfrms']
    req_msg.req.timeout = 10000

    req_msg_raw = req_msg.SerializeToString()

    aesm_service = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    connected = False

    # try to connect to the unnamed socket (for PSW 1.6 and 1.7)
    if not connected:
        try:
            aesm_service.connect("\0sgx_aesm_socket_base" + "\0" * 87)
            connected = True
        except socket.error:
            pass

    # try to connect to the named socket (for PSW 1.8+)
    if not connected:
        try:
            aesm_service.connect("/var/run/aesmd/aesm.socket")
            connected = True
        except socket.error:
            pass

    if not connected:
        raise socket.error("Cannot connect to the AESMD service")

    aesm_service.send(struct.pack("<I", len(req_msg_raw)))
    aesm_service.send(req_msg_raw)

    ret_msg_size = struct.unpack("<I", aesm_service.recv(4))[0]
    ret_msg = aesm_pb2.GetTokenRet()
    ret_msg_raw = aesm_service.recv(ret_msg_size)
    ret_msg.ParseFromString(ret_msg_raw)

    if ret_msg.ret.error != 0:
        raise Exception("Failed. (Error Code = %d)" % (ret_msg.ret.error))

    return ret_msg.ret.token


argparser = argparse.ArgumentParser()
argparser.add_argument('--sig', '-sig', metavar='SIGNATURE',
                       type=argparse.FileType('rb'), required=True,
                       help='Input .sig file (contains SIGSTRUCT)')
argparser.add_argument('--output', '-output', metavar='OUTPUT',
                       type=argparse.FileType('wb'), required=True,
                       help='Output .token file (contains EINITTOKEN)')


def main(args=None):
    """Main Program."""
    args = argparser.parse_args(args)

    attr = read_sigstruct(args.sig.read())

    print("Attributes:")
    print("    mrenclave: %s" % attr['mrenclave'].hex())
    print("    isvprodid: %d" % attr['isvprodid'])
    print("    isvsvn:    %d" % attr['isvsvn'])
    print("    flags:     %016x" % int.from_bytes(attr['flags'],
                                                  byteorder='big'))
    print("    xfrms:     %016x" % int.from_bytes(attr['xfrms'],
                                                  byteorder='big'))
    print("    miscs:     %08x" % int.from_bytes(attr['miscs'],
                                                 byteorder='big'))
    print("    miscmask:  %08x" % int.from_bytes(attr['miscmask'],
                                                 byteorder='big'))
    print("    modulus:   %s..." % attr['modulus'].hex()[:32])
    print("    exponent:  %d" % attr['exponent'])
    print("    signature: %s..." % attr['signature'].hex()[:32])

    token = connect_aesmd(attr)
    args.output.write(token)
    return 0


if __name__ == "__main__":
    sys.exit(main())