tnet/src/util.py

import socket
import selectors

from protocol import DefaultHandler
from protocol import NoReply


def accept_connection(sel, sock, callback):
    """
    Accepts an incoming connection using protocol.DefaultHandler.

    Args:
       sel (selector): A selector from selectors module.
       sock (socket): A socket from socket module.
       callback (Pattern): Implementation of protocol.Pattern.

    Returns:
        None.
    """
    conn, addr = sock.accept()
    tag = addr
    if conn.family == socket.AF_UNIX:
        tag = f'<UNIX-{conn.getsockname()}>'
    conn.setblocking(False)
    handler = DefaultHandler(sel, conn, addr, tag)
    handler.setpattern(callback)
    events = selectors.EVENT_READ
    sel.register(conn, events, data=handler)


def register_udp_handler(sel, ip, port):
    """
    UPD listening with a predefined protocol.DefaultHandler.

    Args:
       sel (selector): A selector from selectors module.
       ip (str): Broadcasting lan IP address.
       port (int): Port to listen to.

    Returns:
        DefaultHandler: The handler registered in selector.
    """
    udp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
                             socket.IPPROTO_UDP)
    udp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    udp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    udp_sock.bind((ip, port))  # no need to listen for UDP
    print(f'Listening on {(ip, port)}')
    udp_sock.setblocking(False)
    udp_handler = DefaultHandler(sel, udp_sock, ip, 'udp')
    # set data=udp_handler right here as accept_connection()
    # cannot be applied in case of UDP
    sel.register(udp_sock, selectors.EVENT_READ, data=udp_handler)
    return udp_handler


def register_unix_handler(sel, path):
    """
    UNIX listening. A handler is to be set on accepting a new connection.

    Args:
       sel (selector): A selector from selectors module.
       path (str): A path to the UNIX socket.

    Returns:
        None.
    """
    unix_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    unix_sock.bind(path)
    unix_sock.listen(1)
    unix_sock.setblocking(False)
    sel.register(unix_sock, selectors.EVENT_READ, data=None)


def register_tcp_handler(sel, ip, port):
    """
    TCP listening. A handler is to be set on accepting a new connection.

    Args:
       sel (selector): A selector from selectors module.
       ip (str): This machine's lan IP address.
       port (int): Port to listen to.

    Returns:
        None.
    """
    tcp_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    tcp_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    tcp_sock.bind((ip, port))
    tcp_sock.listen()
    print(f'Listening on {(ip, port)}')
    tcp_sock.setblocking(False)
    sel.register(tcp_sock, selectors.EVENT_READ, data=None)


def send_via_tcp(sel, ip, port, pattern, *missives):
    """
    Initialize outbound TCP connection with protocol.DefaultHandler.

    Args:
       sel (selector): A selector from selectors module.
       ip (str): This machine's lan IP address.
       port (int): Port to listen to.
       pattern (Pattern): Implementation of protocol.Pattern.
       missives (Missive): A sequence of messages to send.

    Returns:
        None.
    """
    if not missives:
        return
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setblocking(True)
    sock.connect((ip, port))
    events = selectors.EVENT_WRITE
    handler = DefaultHandler(sel, sock, ip, 'tcp-out')
    handler.setpattern(pattern)
    sel.register(sock, events, data=handler)
    for missive in missives:
        handler.enqueue(missive)


def broadcast_via_udp(sel, ip, port, missive):
    """
    Initialize UDP broadcast with protocol.DefaultHandler.

    Args:
       sel (selector): A selector from selectors module.
       ip (str): Broacasting lan IP address.
       port (int): Port for broadcasting.
       missive (Missive): A message to send out.

    Returns:
        None.
    """
    if not missive:
        return
    sock = socket.socket(socket.AF_INET,
                         socket.SOCK_DGRAM,
                         socket.IPPROTO_UDP)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.setblocking(True)
    sock.bind((ip, port))
    events = selectors.EVENT_WRITE
    handler = DefaultHandler(sel, sock, ip, 'udp-out')
    handler.setpattern(NoReply())
    sel.register(sock, events, data=handler)
    handler.enqueue(missive)


def send_via_unix(sel, path, pattern, *missives):
    """
    Initialize UNIX socket connection with protocol.DefaultHandler.

    Args:
       sel (selector): A selector from selectors module.
       path (str): A path to a UNIX socket file.
       pattern (Pattern): Implementation of protocol.Pattern.
       missives (Missive): A sequence of messages to send.

    Returns:
        None.
    """
    if not missives:
        return
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.setblocking(True)
    sock.connect(path)
    events = selectors.EVENT_WRITE
    handler = DefaultHandler(sel, sock,
                             path, 'unix-out')
    handler.setpattern(pattern)
    sel.register(sock, events, data=handler)
    for missive in missives:
        handler.enqueue(missive)