workflow-flash-with-uboot/deployment/serial/init_test.py

"""Restart RUTBS from U-Boot and verify the ordered boot and kernel name."""
import argparse
import fcntl
import os
import re
import time
from pathlib import Path

import serial


STAGES = (
    ('U-Boot prompt', '=>', b'res\r'),
    ('kernel', 'Starting kernel ...', None),
    ('debug selection', 'Press the [1], [2], [3] or [4] key and hit [enter] to select the debug level', None),
    ('console activation', 'Please press Enter to activate this console.', b'\r'),
    ('login', 'RUTBS login:', b'root\r'),
    ('password', 'Password:', 'password'),
    ('shell', 'root@RUTBS:~#', b'uname -s\r'),
    ('command result', 'root@RUTBS:~#', None),
)
FAILURE = re.compile(r'Kernel panic|Login incorrect|Login failed|Authentication fail(?:ed|ure)|Access denied', re.I)


def run_test(port, password, transcript, deadline):
    stage = 0
    pending = ''
    observed = ''
    # Request a fresh prompt when the existing => predates opening the port.
    remaining = deadline - time.monotonic()
    if remaining <= 0:
        raise RuntimeError('Overall test timeout')
    port.write_timeout = min(2, remaining)
    if port.write(b'\r') != 1:
        raise RuntimeError('Incomplete serial write')
    while time.monotonic() < deadline:
        data = port.read(min(port.in_waiting or 1, 4096))
        if not data:
            continue
        transcript.extend(data)
        chunk = data.decode('utf-8', errors='replace').replace('\r', '')
        pending += chunk
        observed = (observed + chunk)[-8192:]
        if FAILURE.search(observed):
            raise RuntimeError('Boot or authentication failure detected')
        while stage < len(STAGES):
            name, token, action = STAGES[stage]
            if stage == 0:
                match = re.search(r'(?:^|\n)[^\n]*=> *$', pending)
                position = pending.rfind('=>') if match else -1
            else:
                position = pending.find(token)
            if position < 0:
                break
            before = pending[:position]
            pending = pending[position + len(token):]
            if stage == len(STAGES) - 1:
                lines = [line.strip() for line in before.splitlines() if line.strip()]
                if lines and lines[0] == 'uname -s':
                    lines.pop(0)  # A terminal may echo the command.
                if lines != ['Linux']:
                    raise RuntimeError('uname -s output was not exactly Linux')
                if time.monotonic() >= deadline:
                    raise RuntimeError('Overall test timeout')
                return
            if action is not None:
                # Discard pre-existing input: the next prompt must follow our write.
                pending = ''
                payload = password.encode('utf-8') + b'\r' if action == 'password' else action
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise RuntimeError('Overall test timeout')
                port.write_timeout = min(2, remaining)
                if port.write(payload) != len(payload):
                    raise RuntimeError('Incomplete serial write')
            stage += 1
        if len(pending) > 65536:
            raise RuntimeError('Unexpected excessive serial output in current stage')
    raise RuntimeError('Timeout waiting for ' + STAGES[stage][0])


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--port', default='/dev/ttyUSB0')
    args = parser.parse_args()
    deadline = time.monotonic() + 180
    password = os.environ.get('RUTBS_PASSWORD', '')
    transcript = bytearray()
    result = 1
    reason = 'Serial connection or configuration error'
    log = Path('artifacts/boot.log')
    try:
        log.parent.mkdir(parents=True, exist_ok=True)
        # Never put raw serial bytes on disk or stdout, even during authentication.
        log.write_text('', encoding='utf-8')
        if not password or any(c in password for c in '\r\n\x00'):
            raise RuntimeError('RUTBS_PASSWORD is missing or invalid')
        with open('/tmp/uboot-serial.lock', 'w') as lock:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            with serial.Serial(args.port, 115200, timeout=0.1, write_timeout=2, exclusive=True) as port:
                run_test(port, password, transcript, deadline)
        result = 0
    except RuntimeError as exc:
        reason = str(exc)
    except Exception:
        # Serial/library exceptions can include transmitted data; do not print them.
        pass
    finally:
        content = transcript.decode('utf-8', errors='replace').replace('\r', '')
        if password:
            content = content.replace(password, '[REDACTED]')
            reason = reason.replace(password, '[REDACTED]')
        try:
            log.write_text(content, encoding='utf-8')
        except OSError:
            result = 1
            reason = 'Could not save serial log'
    if result:
        print('BOOT TEST: FAIL')
        print(reason)
    else:
        print('BOOT TEST: PASS')
    return result


if __name__ == '__main__':
    raise SystemExit(main())