#!/usr/bin/env python

import io
import socket
import struct
import tpm2
from argparse import ArgumentParser
from argparse import FileType
from fcntl import ioctl
from ctypes import Structure
from ctypes import c_uint32
from ctypes import addressof

VTPM_FLAG_TPM2 = 1
VTPM_IOC_NEW_DEV = 0x4014a100


class c_vtpm_new_dev(Structure):
    _fields_ = \
        [('flags', c_uint32),
         ('tpm_num', c_uint32),
         ('fd', c_uint32),
         ('major', c_uint32),
         ('minor', c_uint32)]


def main():
    parser = ArgumentParser(description='Run a TPM 2.0 simulator proxy')
    parser.add_argument('--host', dest='host', metavar='NAMEALG',
                        help='Address of the simulator',
                        type=str, default='localhost')
    args = parser.parse_args()

    sim = tpm2.Simulator(args.host)

    client = tpm2.Client(flags = tpm2.Client.FLAG_DEBUG, simulator = sim)
    startup_cmd = struct.pack('B' * 12,
        0x80, 0x01,
        0x00, 0x00, 0x00, 0x0C,
        0x00, 0x00, 0x01, 0x44,
        0x00, 0x00)
    client.send_cmd(startup_cmd)

    new_dev = c_vtpm_new_dev(flags = VTPM_FLAG_TPM2)

    with open("/dev/vtpmx", "rb+") as vtpmx:
        ioctl(vtpmx, VTPM_IOC_NEW_DEV, addressof(new_dev))

        fp = io.open(new_dev.fd, 'rb+', buffering = 0)

        while True:
            stream = fp.read(4096)
            try:
                resp = client.send_cmd(stream)
            except tpm2.ProtocolError, e:
                print str(e)
            fp.write(resp)


if __name__ == '__main__':
    main()