#!/usr/bin/env python3 # SPDX-License-Identifier: BSD-3-Clause OR Apache-2.0 # encoding: utf-8 import abc import logging import random import socket import uuid from .aap_message import AAPMessage, AAPMessageType, InsufficientAAPDataError logger = logging.getLogger(__name__) class AAPClient(abc.ABC): """A context manager class for connecting to the AAP socket of a uD3TN instance. Args: socket: A `socket.socket` object address: The address of the remote socket to be used when calling `socket.connect()` """ def __init__(self, address): self.socket = None self.address = address self.node_eid = None self.agent_id = None @abc.abstractmethod def connect(self): """Establish a socket connection to a uD3TN instance and return the received welcome message. """ logger.debug("Connected to uD3TN, awaiting WELCOME message...") return self._welcome() def disconnect(self): """Shutdown and close the socket.""" logger.debug("Terminating connection...") self.socket.shutdown(socket.SHUT_RDWR) self.socket.close() def __enter__(self): """Return `self` upon calling `self.connect()` to establish the socket connection. """ self.connect() return self def __exit__(self, type, value, traceback): """Invoke `self.disconnect()` on any raised runtime exception.""" self.disconnect() def _welcome(self): """Receive the AAP welcome message and store the node EID of the uD3TN instance in `self.node_eid`. """ msg_welcome = self.receive() if msg_welcome.msg_type != AAPMessageType.WELCOME: raise RuntimeError( f"Expected WELCOME message, got {msg_welcome.msg_type}" ) logger.debug(f"WELCOME message received! ~ EID = {msg_welcome.eid}") self.node_eid = msg_welcome.eid return msg_welcome @property def eid(self): """Return the own EID.""" if self.node_eid[0:3] == "ipn": prefix, _ = self.node_eid.split(".") return f"{prefix}.{self.agent_id}" else: return f"{self.node_eid}{self.agent_id}" def _generate_eid(self): if not self.node_eid: return None if self.node_eid[0:3] == "dtn": return str(uuid.uuid4()) elif self.node_eid[0:3] == "ipn": return str( random.randint(1, 4294967295) ) # UINT32_MAX for maximum compatibility return None def register(self, agent_id=None): """Attempt to register the specified agent identifier. Args: agent_id: The agent identifier to be registered. If None, uuid.uuid4() is called to generate one. """ self.agent_id = agent_id or self._generate_eid() logger.debug(f"Sending REGISTER message for '{agent_id}'...") msg_ack = self.send( AAPMessage(AAPMessageType.REGISTER, self.agent_id) ) if msg_ack.msg_type != AAPMessageType.ACK: raise RuntimeError(f"Expected ACK message, got {msg_ack.msg_type}") logger.debug("ACK message received!") def ping(self): """Send a PING message via AAP and returns the ACK message (e.g. for keepalive purposes). """ msg_ack = self.send(AAPMessage(AAPMessageType.PING)) if msg_ack.msg_type != AAPMessageType.ACK: raise RuntimeError(f"Expected ACK message, got {msg_ack.msg_type}") return msg_ack def receive(self): """Receive and return the next `AAPMessage`.""" buf = bytearray() msg = None while msg is None: data = self.socket.recv(1) if not data: logger.debug("Disconnected") return None buf += data try: msg = AAPMessage.parse(buf) except InsufficientAAPDataError: continue return msg def send(self, aap_msg): """Serialize and send the provided `AAPMessage` to the AAP endpoint. Args: aap_msg: The `AAPMessage` to be sent. """ self.socket.send(aap_msg.serialize()) return self.receive() def send_bundle(self, dest_eid, bundle_data, bibe=False): """Send the provided bundle to the AAP endpoint. Args: dest_eid: The destination EID. bundle_data: The binary payload data to be encapsulated in a bundle. bibe: Whether the AAPClient should send a SENDBIBE message (True) or a SENDBUNDLE message (False). Defaults to False. """ if not bibe: logger.debug(f"Sending SENDBUNDLE message to {dest_eid}") msg_sendconfirm = self.send(AAPMessage( AAPMessageType.SENDBUNDLE, dest_eid, bundle_data )) if msg_sendconfirm.msg_type != AAPMessageType.SENDCONFIRM: raise RuntimeError( f"Expected SENDCONFIRM message, " f"got {msg_sendconfirm.msg_type}" ) try: bundle_id = msg_sendconfirm.decode_bundle_id() except ValueError: bundle_id = msg_sendconfirm.bundle_id logger.debug( f"SENDCONFIRM message received! ~ ID = {bundle_id}" ) return msg_sendconfirm else: logger.debug(f"Sending SENDBIBE message to {dest_eid}") msg_sendconfirm = self.send(AAPMessage( AAPMessageType.SENDBIBE, dest_eid, bundle_data )) if msg_sendconfirm.msg_type != AAPMessageType.SENDCONFIRM: raise RuntimeError( f"Expected SENDCONFIRM message, " f"got {msg_sendconfirm.msg_type}" ) try: bundle_id = msg_sendconfirm.decode_bundle_id() except ValueError: bundle_id = msg_sendconfirm.bundle_id logger.debug( f"SENDCONFIRM message received! ~ ID = {bundle_id}" ) return msg_sendconfirm def send_str(self, dest_eid, bundle_data): """Send the provided bundle to the AAP endpoint. Args: dest_eid: The destination EID. bundle_data: The string message to be utf-8 encoded and encapsulated in a bundle. """ return self.send_bundle(dest_eid, bundle_data.encode("utf-8"), False) class AAPUnixClient(AAPClient): """A context manager class for connecting to the AAP Unix socket of a uD3TN instance. Args: address: The address (PATH) of the remote socket to be used when calling `socket.connect()` """ def __init__(self, address='ud3tn.socket'): super().__init__(address=address) def connect(self): self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) self.socket.connect(self.address) return super().connect() class AAPTCPClient(AAPClient): """A context manager class for connecting to the AAP TCP socket of a uD3TN instance. Args: address: The address tupel (HOST, PORT) of the remote socket to be used when calling `socket.connect()` """ def __init__(self, address=('localhost', '4242')): super().__init__(address=address) def connect(self): host, service = self.address addrinfo = socket.getaddrinfo(host, service, socket.AF_UNSPEC, socket.SOCK_STREAM) for af, socktype, proto, canonname, sa in addrinfo: try: s = socket.socket(af, socktype, proto) except OSError: s = None continue try: s.connect(sa) except OSError: s.close() s = None continue break if s is None: raise RuntimeError(f"cannot connect to {host}:{service}") self.socket = s return super().connect()