import os
import shutil
import string
import subprocess

def strip_special(line):
    return ''.join([c for c in line if c in string.printable])

ovmf_vars = '/usr/share/edk2/ovmf/OVMF_VARS.fd'
ovmf_binary = '/usr/share/edk2/ovmf/OVMF_CODE.secboot.fd'
uefi_shell = '/usr/share/edk2/ovmf/UefiShell.iso'
vars_output = 'output_vars.fd'

if os.path.exists(vars_output):
    raise Exception('%s already exists' % vars_output)

shutil.copy(ovmf_vars, vars_output)

cmd = [
    '/usr/bin/qemu-system-x86_64',
    '-machine', 'q35,smm=on,accel=kvm',
    '-display', 'none',
    '-no-user-config',
    '-nodefaults',
    '-m', '256',
    '-smp', '2,sockets=2,cores=1,threads=1',
    '-chardev', 'pty,id=charserial1',
    '-device', 'isa-serial,chardev=charserial1,id=serial1',
    '-global', 'driver=cfi.pflash01,property=secure,value=on',
    '-drive', 'file=%s,if=pflash,format=raw,unit=0,readonly=on' % ovmf_binary,
    '-drive', 'file=%s,if=pflash,format=raw,unit=1,readonly=off' % vars_output,
    '-drive', 'file=%s,format=raw,if=none,media=cdrom,id=drive-virtio-disk1,readonly=on' % uefi_shell,
    '-device', 'virtio-blk-pci,scsi=off,drive=drive-virtio-disk1,id=virtio-disk1,bootindex=1',
    '-object', 'rng-random,id=objrng0,filename=/dev/urandom',
    '-device', 'virtio-rng-pci,rng=objrng0,id=rng0',
    '-serial', 'stdio']
p = subprocess.Popen(cmd,
    stdin=subprocess.PIPE,
    stdout=subprocess.PIPE,
    stderr=subprocess.PIPE)
# Wait until the shell starts (First line printed)
print('OUT: %s' % strip_special(p.stdout.readline()))
# Send escape to enter shell early
p.stdin.write(chr(27))
p.stdin.write(b'fs0:\r\n')
p.stdin.write(b'EnrollDefaultKeys.efi\r\n')
p.stdin.write(b'reset\r\n')
while True:
    read = p.stdout.readline()
    print('OUT: %s' % strip_special(read))
    if 'info: success' in read:
        break
p.kill()
print(strip_special(p.stdout.read()))
print('Created %s' % vars_output)