# fluff

import atexit
import json
import socket
import subprocess
import time

QMP_PORT = 12345
TEST_PORT = 12346
QEMU = 'qemu-system-x86_64'
ARGS = ['-nographic',
        '-qmp', f'tcp:127.0.0.1:{QMP_PORT},server,nowait,nodelay',
        '-device', 'virtio-net-pci,netdev=net0',
        '-netdev', 'user,id=net0,restrict=on']

p = subprocess.Popen([QEMU] + ARGS, stdout=subprocess.DEVNULL)
atexit.register(p.kill)

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
time.sleep(.2)
try:
    print('QMP connecting, attempt 1/2')
    sock.connect(('127.0.0.1', QMP_PORT))
except ConnectionRefusedError:
    print('QMP connecting, attempt 2/2')
    time.sleep(2)
    sock.connect(('127.0.0.1', QMP_PORT))
print('QMP connected...')


def execute(cmd, **kwargs):
    sock.send(json.dumps(
        {'execute': cmd, 'arguments': kwargs} if kwargs else {'execute': cmd}
    ).encode())
    print('executed:', cmd, kwargs)


def expect(validator_func=None):
    r = b''
    while True:
        r += sock.recv(1)
        if r[-1] == ord('\n'):
            reply = json.loads(r.decode())
            print(f'received: {reply}')
            if validator_func is not None:
                assert validator_func(reply)
            return reply


execute('qmp_capabilities')
expect(lambda r: 'QMP' in r)


# inspect as started up

execute('human-monitor-command', **{'command-line': 'info network'})
expect(lambda r: r == {'return': {}})
reply = expect(lambda r: set(r.keys()) == {'return'})
# Should be:
# virtio-net-pci.0: index=0,type=nic,model=virtio-net-pci,macaddr=...
#  \ net0: index=0,type=user,net=10.0.2.0,restrict=on
assert('restrict=on' in reply['return'])
print()


# modify

execute('netdev_add', type='user', id='net0',
        restrict=False,  # unrestrict
        hostfwd=[{'str': f'tcp:127.0.0.1:{TEST_PORT}-:1'}])
        # ^ *not* a list of strings as the documentation claims


# inspect modified

execute('human-monitor-command', **{'command-line': 'info network'})
expect(lambda r: r == {'return': {}})
reply = expect(lambda r: set(r.keys()) == {'return'})
# And now we have two net0s:
# virtio-net-pci.0: index=0,type=nic,model=virtio-net-pci,macaddr=...
#  \ net0: index=0,type=user,net=10.0.2.0,restrict=on
#  \ net0: index=0,type=user,net=10.0.2.0,restrict=off
assert('restrict=on' in reply['return'])
assert('restrict=off' in reply['return'])


# connect to modified

sock_test = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
time.sleep(.2)
sock_test.connect(('127.0.0.1', TEST_PORT))
sock_test.close()
print(f'port {TEST_PORT} accepting connections')