sedbgmux/handlers.py

# This file is a part of sedbgmux, an open source DebugMux client.
# Copyright (c) 2022  Vadim Yanitskiy <axilirator@gmail.com>
#
# SPDX-License-Identifier: GPL-3.0-or-later
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

import logging as log
import threading
import socket
import queue
import enum
import abc
import sys

from typing import Any, Optional
from construct import Container
from proto import DbgMuxFrame


class DbgMuxFrameHandler(abc.ABC):
    ''' Abstract DebugMux frame handler '''

    _tx_queue: Optional[queue.Queue] = None

    def send(self, msg_type: DbgMuxFrame.MsgType, msg: Any = b''):
        ''' Called by child classes to send some message '''
        assert self._tx_queue is not None
        self._tx_queue.put((msg_type, msg))

    @abc.abstractmethod
    def _handle_frame(self, frame: Container) -> None:
        ''' Handle the given DebugMux frame '''


class ConnState(enum.Enum):
    ''' Connection state for DbgMuxConnHandler '''
    NotEstablished      = enum.auto()
    Establishing        = enum.auto()
    Established         = enum.auto()


class DbgMuxConnHandler(DbgMuxFrameHandler):
    ''' Abstract DebugMux connection handler '''

    def __init__(self):
        self.conn_state: ConnState = ConnState.NotEstablished
        self.ConnRef: int = 0xffff
        self.DPRef: int = 0xffff

    def establish(self, DPRef: int) -> None:
        assert self.conn_state == ConnState.NotEstablished
        log.info("Establishing connection with DPRef=0x%04x", DPRef)
        self.send(DbgMuxFrame.MsgType.ConnEstablish, dict(DPRef=DPRef))
        self.conn_state = ConnState.Establishing
        self.DPRef = DPRef

    def send_data(self, data: bytes) -> None:
        ''' Called by child classes to send connection data '''
        assert self.conn_state == ConnState.Established
        msg = dict(ConnRef=self.ConnRef, Data=data)
        self.send(DbgMuxFrame.MsgType.ConnData, msg)

    def _match(self, frame: Container,
                     msg_type: DbgMuxFrame.MsgType,
                     msg_fields: dict = { }) -> bool:
        if frame['MsgType'] != msg_type:
            return False
        for (key, val) in msg_fields.items():
            if key not in frame['Msg']:
                return False
            if frame['Msg'][key] != val:
                return False
        return True

    def _handle_frame(self, frame: Container) -> None:
        ''' Handle the given DebugMux frame '''

        if self.conn_state == ConnState.Established:
            fields = dict(ConnRef=self.ConnRef)
            if self._match(frame, DbgMuxFrame.MsgType.ConnData, fields):
                self._handle_data(frame['Msg']['Data'])
                self.send(DbgMuxFrame.MsgType.Ack)
                raise StopIteration
            elif self._match(frame, DbgMuxFrame.MsgType.ConnTerminated, fields):
                log.info('Connection terminated (DPRef=0x%04x, ConnRef=0x%04x)',
                         self.DPRef, self.ConnRef)
                self.conn_state = ConnState.NotEstablished
                self._handle_terminate()
                self.send(DbgMuxFrame.MsgType.Ack)
                raise StopIteration
        elif self.conn_state == ConnState.Establishing:
            # Match ConnEstablished with our DPRef and any ConnRef
            if self._match(frame, DbgMuxFrame.MsgType.ConnEstablished, dict(DPRef=self.DPRef)):
                log.info('Connection established (DPRef=0x%04x, ConnRef=0x%04x)',
                         self.DPRef, frame['Msg']['ConnRef'])
                self.conn_state = ConnState.Established
                self.ConnRef = frame['Msg']['ConnRef']
                self._handle_establish()
                self.send(DbgMuxFrame.MsgType.Ack)
                raise StopIteration

    @abc.abstractmethod
    def _handle_data(self, data: bytes) -> None:
        ''' Called on reciept of connection data '''

    @abc.abstractmethod
    def _handle_establish(self) -> None:
        ''' Called on connection establishment '''

    @abc.abstractmethod
    def _handle_terminate(self) -> None:
        ''' Called on connection termination '''


class DbgMuxConnInteractiveTerminal(DbgMuxConnHandler):
    def __init__(self, *args, **kw):
        self.attached: bool = False
        super().__init__(*args)

    def attach(self):
        self.attached = True
        while True:
            try:
                line = input()
                if self.conn_state == ConnState.Established:
                        self.send_data(bytes(line, 'ascii'))
            except (KeyboardInterrupt, EOFError) as e:
                break
        self.attached = False

    def _handle_data(self, data: bytes) -> None:
        if self.attached:
            sys.stdout.write(data.decode('ascii'))

    def _handle_establish(self) -> None:
        pass

    def _handle_terminate(self) -> None:
        pass


class DbgMuxConnUdpBridge(DbgMuxConnHandler):
    DGRAM_MAX_LEN: int = 1024

    def __init__(self, *args, **kw):
        super().__init__(*args)

        self.raddr: str = kw.get('raddr', '127.0.0.1')
        self.rport: int = kw.get('rport', 9999)
        self.laddr: str = kw.get('laddr', '127.0.0.1')
        self.lport: int = kw.get('lport', 8888)

        self._sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._sock.bind((self.laddr, self.lport))
        self._sock.settimeout(0.5)

        self._thread = threading.Thread(target=self._worker,
                                        daemon=True)
        self._shutdown = threading.Event()

    def _worker(self) -> None:
        while not self._shutdown.is_set():
            try:
                (data, addr) = self._sock.recvfrom(self.DGRAM_MAX_LEN)
                if self.conn_state == ConnState.Established:
                    self.send_data(data)
            except TimeoutError:
                pass

    def _handle_data(self, data: bytes) -> None:
        self._sock.sendto(data, (self.raddr, self.rport))

    def _handle_establish(self) -> None:
        self._shutdown.clear()
        self._thread.start()

    def _handle_terminate(self) -> None:
        self._shutdown.set()