|
| 1 | +#!/usr/bin/python3 |
| 2 | +import fcntl |
| 3 | +import logging |
| 4 | +import os |
| 5 | +import signal |
| 6 | +import socket |
| 7 | +import struct |
| 8 | +import subprocess |
| 9 | +import sys |
| 10 | +import telnetlib |
| 11 | +import time |
| 12 | + |
| 13 | +BASE_IMG = '/sonic-vs.img' |
| 14 | +USER = 'admin' |
| 15 | +PASSWORD = 'YourPaSsWoRd' |
| 16 | + |
| 17 | + |
| 18 | +class Qemu: |
| 19 | + def __init__(self, name: str, memory: str, interfaces: int): |
| 20 | + self._name = name |
| 21 | + self._memory = memory |
| 22 | + self._interfaces = interfaces |
| 23 | + self._p = None |
| 24 | + self._disk = '/overlay.img' |
| 25 | + |
| 26 | + def prepare_overlay(self, base: str) -> None: |
| 27 | + cmd = [ |
| 28 | + 'qemu-img', |
| 29 | + 'create', |
| 30 | + '-f', 'qcow2', |
| 31 | + '-b', base, |
| 32 | + self._disk, |
| 33 | + ] |
| 34 | + subprocess.run(cmd, check=True) |
| 35 | + |
| 36 | + def start(self) -> None: |
| 37 | + cmd = [ |
| 38 | + 'qemu-system-x86_64', |
| 39 | + '-cpu', 'host', |
| 40 | + '-display', 'none', |
| 41 | + '-enable-kvm', |
| 42 | + '-machine', 'q35', |
| 43 | + '-name', self._name, |
| 44 | + '-m', self._memory, |
| 45 | + '-drive', f'if=virtio,format=qcow2,file={self._disk}', |
| 46 | + '-serial', 'telnet:127.0.0.1:5000,server,nowait', |
| 47 | + ] |
| 48 | + |
| 49 | + for i in range(self._interfaces): |
| 50 | + with open(f'/sys/class/net/eth{i}/address', 'r') as f: |
| 51 | + mac = f.read().strip() |
| 52 | + cmd.append('-device') |
| 53 | + cmd.append(f'virtio-net,netdev=hn{i},mac={mac}') |
| 54 | + cmd.append(f'-netdev') |
| 55 | + cmd.append(f'tap,id=hn{i},ifname=tap{i},script=/mini-lab/mirror_tap_to_eth.sh,downscript=no') |
| 56 | + |
| 57 | + self._p = subprocess.Popen(cmd) |
| 58 | + |
| 59 | + def wait(self) -> None: |
| 60 | + self._p.wait() |
| 61 | + |
| 62 | + |
| 63 | +class Telnet: |
| 64 | + def __init__(self): |
| 65 | + self._tn: telnetlib.Telnet | None = None |
| 66 | + |
| 67 | + def connect(self, host: str, port: int, max_retries=60) -> bool: |
| 68 | + for i in range(1, max_retries + 1): |
| 69 | + try: |
| 70 | + self._tn = telnetlib.Telnet(host, port) |
| 71 | + return True |
| 72 | + except: |
| 73 | + time.sleep(1) |
| 74 | + if i == max_retries: |
| 75 | + return False |
| 76 | + |
| 77 | + def close(self): |
| 78 | + self._tn.close() |
| 79 | + |
| 80 | + def wait_until(self, match: str): |
| 81 | + self._tn.read_until(match.encode('ascii')) |
| 82 | + |
| 83 | + def write_and_wait(self, data: str, match: str = '$ ') -> str: |
| 84 | + self._tn.write(data.encode('ascii') + b'\n') |
| 85 | + return self._tn.read_until(match.encode('ascii')).decode('utf-8') |
| 86 | + |
| 87 | + def write_test(self, data: str) -> str: |
| 88 | + self._tn.write(data.encode('ascii') + b'\n') |
| 89 | + time.sleep(5) |
| 90 | + return self._tn.read_some().decode('utf-8') |
| 91 | + |
| 92 | + |
| 93 | +def main(): |
| 94 | + signal.signal(signal.SIGINT, handle_exit) |
| 95 | + signal.signal(signal.SIGTERM, handle_exit) |
| 96 | + |
| 97 | + logging.basicConfig(level=logging.INFO, stream=sys.stdout) |
| 98 | + logger = logging.getLogger() |
| 99 | + |
| 100 | + name = os.getenv('CLAB_LABEL_CLAB_NODE_NAME', default='switch') |
| 101 | + memory = os.getenv('VM_MEMORY', default='2048') |
| 102 | + interfaces = int(os.getenv('CLAB_INTFS', 0)) + 1 |
| 103 | + |
| 104 | + logger.info(f'Waiting for {interfaces} interfaces to be connected') |
| 105 | + wait_until_all_interfaces_are_connected(interfaces) |
| 106 | + |
| 107 | + vm = Qemu(name, memory, interfaces) |
| 108 | + |
| 109 | + logger.info('Prepare disk') |
| 110 | + vm.prepare_overlay(BASE_IMG) |
| 111 | + |
| 112 | + logger.info('Start QEMU') |
| 113 | + vm.start() |
| 114 | + |
| 115 | + logger.info('Try to connect via telnet...') |
| 116 | + tn = Telnet() |
| 117 | + if not tn.connect('127.0.0.1', 5000): |
| 118 | + logger.error('Cannot connect to telnet server') |
| 119 | + sys.exit(1) |
| 120 | + |
| 121 | + logger.info('Connected via telnet and waiting for login prompt') |
| 122 | + tn.wait_until('login: ') |
| 123 | + |
| 124 | + logger.info('Try to login') |
| 125 | + tn.write_and_wait(USER, 'Password: ') |
| 126 | + tn.write_and_wait(PASSWORD) |
| 127 | + |
| 128 | + logger.info('Authorize ssh key') |
| 129 | + authorize_ssh_key(tn) |
| 130 | + |
| 131 | + logger.info('Wait until config-setup is done') |
| 132 | + if not wait_until_config_setup_is_done(tn): |
| 133 | + logger.error('config-setup still not done') |
| 134 | + sys.exit(1) |
| 135 | + |
| 136 | + net = get_ip_address('eth0') + '/16' |
| 137 | + logger.info(f'Configure {net} on eth0') |
| 138 | + tn.write_and_wait(f'sudo config interface ip add eth0 {net}') |
| 139 | + tn.write_and_wait('sudo config save --yes') |
| 140 | + |
| 141 | + tn.close() |
| 142 | + |
| 143 | + logger.info('Wait until QEMU terminated') |
| 144 | + vm.wait() |
| 145 | + |
| 146 | + |
| 147 | +def handle_exit(signal, frame): |
| 148 | + sys.exit(0) |
| 149 | + |
| 150 | + |
| 151 | +def wait_until_all_interfaces_are_connected(interfaces: int) -> None: |
| 152 | + while True: |
| 153 | + i = 0 |
| 154 | + for iface in os.listdir('/sys/class/net/'): |
| 155 | + if iface.startswith('eth'): |
| 156 | + i += 1 |
| 157 | + if i == interfaces: |
| 158 | + break |
| 159 | + time.sleep(1) |
| 160 | + |
| 161 | + |
| 162 | +def get_ip_address(iface: str) -> str: |
| 163 | + # Source: https://bit.ly/3dROGBN |
| 164 | + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) |
| 165 | + return socket.inet_ntoa(fcntl.ioctl( |
| 166 | + s.fileno(), |
| 167 | + 0x8915, # SIOCGIFADDR |
| 168 | + struct.pack('256s', iface.encode('utf_8')) |
| 169 | + )[20:24]) |
| 170 | + |
| 171 | + |
| 172 | +def wait_until_config_setup_is_done(tn: Telnet, max_retries: int = 60) -> bool: |
| 173 | + for i in range(1, max_retries + 1): |
| 174 | + # updategraph is started after the config-setup |
| 175 | + result = tn.write_and_wait('systemctl is-active updategraph') |
| 176 | + if not 'inactive' in result: |
| 177 | + return True |
| 178 | + time.sleep(1) |
| 179 | + if i == max_retries: |
| 180 | + return False |
| 181 | + |
| 182 | + |
| 183 | +def authorize_ssh_key(tn: Telnet) -> None: |
| 184 | + with open('/id_rsa.pub') as f: |
| 185 | + key = f.read().strip() |
| 186 | + |
| 187 | + tn.write_and_wait(f'echo "{key}" > authorized_keys') |
| 188 | + tn.write_and_wait('sudo mkdir /root/.ssh') |
| 189 | + tn.write_and_wait('sudo chmod 0600 /root/.ssh') |
| 190 | + tn.write_and_wait('sudo cp authorized_keys /root/.ssh/') |
| 191 | + |
| 192 | + |
| 193 | +if __name__ == '__main__': |
| 194 | + main() |
0 commit comments