#!/usr/bin/env python

from __future__ import print_function
import argparse
import pexpect
import serial
import sys
import time
from pexpect import fdpexpect

parser = argparse.ArgumentParser(description='Connect to serial console ' +
                                 'to execute stuff')
parser.add_argument('--port', required=True,
                    help='serial console device to connect ' +
                    'to (e.g. /dev/pts/X)')
parser.add_argument('--hostname', default="stretch",
                    help='hostname of the system for login process ' +
                    '(default: stretch)')
parser.add_argument('--pipefile', default="./serial0",
                    help='file name for named pipe file (for ' +
                    'interacting between host + VM via QEMU ' +
                    '(default: ./serial0)')
parser.add_argument('--vmoutput', default="vm-output.log",
                    help='filename for VM output (default: vm-output.log)')
parser.add_argument('--user', default="root",
                    help='user name to use for login (default: root)')
parser.add_argument('--password', default="grml",
                    help='password for login (default: grml)')
args = parser.parse_args()


def execute(port, hostname, user, baudrate=115200, timeout=5):
    ser = serial.Serial(port, baudrate)
    ser.flushInput()
    ser.flushOutput()
    ser.write("\n")
    ser.flush()

    child = fdpexpect.fdspawn(ser.fileno())
    child.sendline("")
    try:
        print("Begin of execution inside VM")
        child.expect("%s@%s" % (user, hostname), timeout=timeout)
        child.sendline("/usr/local/bin/goss --gossfile /root/goss.yaml " +
                       "validate --format tap > /root/goss.tap ; " +
                       "echo $? > /root/goss.exitcode\n")
        # NOTE - the serial0 is hardcoded here
        child.sendline("cat /root/goss.tap > /dev/virtio-ports/serial0\n")
        child.sendline("grep -q '^0' /root/goss.exitcode && " +
                       "echo clean_exit > /dev/virtio-ports/serial0\n")
        child.sendline("grep -q '^0' /root/goss.exitcode || " +
                       "echo failure_exit > /dev/virtio-ports/serial0\n")
        child.sendline("poweroff\n")
        print("End of execution inside VM")
    except Exception as except_inst:
        print("Execution inside VM failed: ", except_inst)


def login(port, hostname, user, password,
          baudrate=115200, timeout=5):
    ser = serial.Serial(port, baudrate)
    ser.flushInput()
    ser.flushOutput()

    child = fdpexpect.fdspawn(ser.fileno())
    child.sendline("\n")

    try:
        child.expect("root@%s" % hostname, timeout=timeout)
        return
    except:
        pass

    print("Checking for login prompt...")
    child.expect("%s login:" % hostname, timeout=timeout)
    ser.write("%s\n" % user)
    ser.flush()
    time.sleep(1)
    ser.write("%s\n" % password)
    ser.flush()
    time.sleep(1)
    print("login ok...")


if __name__ == "__main__":
    hostname = args.hostname
    password = args.password
    pipefile = args.pipefile
    port = args.port
    user = args.user
    vmoutput = args.vmoutput

    with open(pipefile, 'r') as output_pipe:
        success = False
        for i in range(12):
            try:
                print("Logging into {0} via serial "
                      "console [try {1}]".format(port, i))
                login(port, hostname, user, password)
                success = True
                break
            except Exception as except_inst:
                print("Login failure (try {0}):".format(i),
                      except_inst, file=sys.stderr)
                time.sleep(5)

        if success:
            execute(port, hostname, user)
            with open(vmoutput, 'w') as fp:
                output = output_pipe.read()
                print(output)
                fp.write(output)