Refactor: Improve netbyte core functionality and testing

This commit introduces several improvements to the netbyte tool:

- **Core Module:** A new `netbyte/core.py` module is created to house reusable logic, including `to_hex` and `parse_hex_bytes` functions. This promotes better code organization and testability.
- **Enhanced `to_hex`:** The `to_hex` function now correctly handles byte strings and provides more robust ASCII references.
- **New `parse_hex_bytes`:** A new function `parse_hex_bytes` is added to parse various hex string formats, making it easier to send hex data.
- **Improved Error Handling:** `print_error` now raises `SystemExit` for cleaner exit handling.
- **Listen Mode:** Added `--listen` and `--bind` arguments to enable listening for incoming connections (TCP/UDP).
- **Hex Input:** Introduced `--send-hex` to interpret stdin as hex bytes before sending.
- **Test Suite:** A basic test suite (`tests/test_core.py`) is added to verify the core functionality.
- **CI Workflow:** A GitHub Actions workflow (`.github/workflows/tests.yml`) is added to automatically run tests on pushes and pull requests.
- **Dependencies:** Updated `setup.py` to include `pytest` for testing and specify Python 3.9+ compatibility.
- **Test Server:** Minor updates to `testserver.py` for Python 3 compatibility and byte string handling.

Co-authored-by: henry <henry@prince.co>

Author: cursoragent

.github/workflows/tests.yml

@@ -0,0 +1,20 @@
+name: tests
+
+on:
+  push:
+  pull_request:
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - name: Install
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e '.[test]'
+      - name: Run tests
+        run: pytest -q

netbyte/core.py

@@ -0,0 +1,102 @@
+"""Core helpers for netbyte.
+
+This module is intentionally free of socket/CLI code so it can be unit tested.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+_SYMBOLS = set("~`!@#$%^&*()_-+={}[]:>;',</?*-+")
+
+
+def is_symbol(ch: str) -> bool:
+    return ch in _SYMBOLS
+
+
+def _is_printable_reference_byte(b: int) -> bool:
+    # Keep the original tool's intent: annotate common readable bytes.
+    if b < 0x20 or b >= 0x7F:
+        return False
+    ch = chr(b)
+    return ch.isalpha() or ch.isdigit() or is_symbol(ch)
+
+
+def to_hex(data: bytes) -> str:
+    """Convert bytes to formatted hex with lightweight ASCII references.
+
+    Output is similar to the original Python 2 implementation, e.g.:
+        48(H) 65(e) 6C(l) 6C(l) 6F(o) 0D 0A(\n)
+
+    Newlines (0x0A) are rendered with an explicit (\n) marker and a real newline.
+    """
+
+    results: list[str] = []
+    new_line = True
+
+    for b in data:
+        hex_value = f"{b:02X}"
+
+        if _is_printable_reference_byte(b):
+            hex_value = f"{hex_value}({chr(b)})"
+
+        if not new_line:
+            hex_value = " " + hex_value
+
+        if b == 0x0A and not data.isspace():
+            hex_value = hex_value + "(\\n)\n"
+            new_line = True
+        else:
+            new_line = False
+
+        results.append(hex_value)
+
+    return "".join(results)
+
+
+_HEX_PAIR = set("0123456789abcdefABCDEF")
+
+
+def parse_hex_bytes(text: str) -> bytes:
+    """Parse a user-supplied hex string into bytes.
+
+    Accepts common separators/spaces and optional 0x prefixes.
+
+    Examples:
+        "DE AD BE EF" -> b"\xDE\xAD\xBE\xEF"
+        "0xDE,0xAD"   -> b"\xDE\xAD"
+        "deadbeef"    -> b"\xDE\xAD\xBE\xEF"
+    """
+
+    # Strip 0x prefixes and keep only hex digits.
+    cleaned_chars: list[str] = []
+    i = 0
+    while i < len(text):
+        if text[i : i + 2].lower() == "0x":
+            i += 2
+            continue
+        ch = text[i]
+        if ch in _HEX_PAIR:
+            cleaned_chars.append(ch)
+        i += 1
+
+    cleaned = "".join(cleaned_chars)
+
+    if len(cleaned) == 0:
+        return b""
+
+    if len(cleaned) % 2 != 0:
+        raise ValueError("hex input must contain an even number of hex digits")
+
+    return bytes(int(cleaned[i : i + 2], 16) for i in range(0, len(cleaned), 2))
+
+
+@dataclass(frozen=True)
+class NetbyteConfig:
+    udp: bool = False
+    timeout: float = 2.0
+    display_encoding: str = "utf-8"
+    display_errors: str = "replace"
+    send_encoding: str = "utf-8"
+    send_errors: str = "strict"

netbyte/netbyte.py

@@ -13,9 +13,11 @@
 import argparse
 import sys
 import time
-from colorama import Fore, Style
+from colorama import Fore, Style, init as colorama_init
 from threading import Thread
-from Queue import Queue, Empty
+from queue import Queue, Empty
+
+from netbyte.core import to_hex, parse_hex_bytes
 
 
 def is_symbol(character):
@@ -32,66 +34,19 @@ def is_symbol(character):
         return True
 
 
-def to_hex(string):
-    '''
-    Converts string to formatted hex and ASCII reference values
-
-    Returns:
-        str: Formatted hex & ASCII reference
-    '''
-
-    results = []
-
-    new_line = True
-
-    for character in string:
-
-        # Convert ASCII to unicode
-        unicode_value = ord(character)
-
-        # Convert unicode to hex
-        hex_value = hex(unicode_value).replace('0x', '')
-
-        # Add a preceding 0
-        if len(hex_value) == 1:
-            hex_value = '0' + hex_value
-
-        # Make upper case
-        hex_value = hex_value.upper()
-
-        # Add reference ASCII for readability
-        if character.isalpha() or character.isdigit() or is_symbol(character):
-            hex_value = hex_value + '(' + character + ')'
-
-        # Add a space in between hex values (not to a new line)
-        if not new_line:
-            hex_value = ' ' + hex_value
-
-        # Add a newline for readability (corresponding to ASCII)
-        if '0A' in hex_value and not string.isspace():
-            hex_value = hex_value + '(\\n)\n'
-            # Next line will be a newline
-            new_line = True
-        else:
-            new_line = False
-
-        results.append(hex_value)
-
-    full_hex = ''
-    for result in results:
-        full_hex += result
-
-    return full_hex
-
-
-def print_ascii(string):
+def print_ascii(data):
     '''
     Print string with ASCII color configuration
     '''
-    if string.isspace():
+    if isinstance(data, (bytes, bytearray)):
+        text = bytes(data).decode("utf-8", errors="replace")
+    else:
+        text = str(data)
+
+    if text.isspace():
         # Add a space to show colors on a non-ASCII line
-        string = ' ' + string
-    print(Fore.MAGENTA + Style.BRIGHT + string + Style.RESET_ALL)
+        text = ' ' + text
+    print(Fore.MAGENTA + Style.BRIGHT + text + Style.RESET_ALL)
 
 
 def print_hex(string):
@@ -106,7 +61,7 @@ def print_error(string):
     Print string with error color configuration and exit with code 1
     '''
     print(Fore.RED + Style.BRIGHT + string + Style.RESET_ALL)
-    exit(1)
+    raise SystemExit(1)
 
 
 class ReadAsync(object):
@@ -156,12 +111,33 @@ def parse_arguments():
     parser.add_argument('hostname', metavar='HOSTNAME', help='Host or IP to connect to')
     parser.add_argument('port', metavar='PORT', help='Connection port')
     parser.add_argument('-u', dest='udp', action="store_true", default=False, help='Use UDP instead of default TCP')
+    parser.add_argument(
+        '-l',
+        '--listen',
+        dest='listen',
+        action='store_true',
+        default=False,
+        help='Listen instead of connect (TCP: accept one client; UDP: receive from a peer)',
+    )
+    parser.add_argument(
+        '--bind',
+        dest='bind',
+        default=None,
+        help='Bind address for --listen (defaults to HOSTNAME)',
+    )
+    parser.add_argument(
+        '--send-hex',
+        dest='send_hex',
+        action='store_true',
+        default=False,
+        help='Interpret stdin as hex bytes before sending (e.g. \"DE AD BE EF\" or \"0xDE,0xAD\")',
+    )
 
     if len(sys.argv) == 1:
 
         parser.print_help()
 
-        exit(1)
+        raise SystemExit(1)
 
     args = parser.parse_args()
     return args
@@ -172,46 +148,107 @@ def main():
     Main function: Connects to host/port and spawns ReadAsync
     '''
 
+    colorama_init()
     args = parse_arguments()
 
-    if args.udp:
-        connection = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-    else:
-        connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    peer = None
 
-    connection.settimeout(2)
+    if args.listen:
+        bind_host = args.bind if args.bind is not None else args.hostname
+        bind_addr = (bind_host, int(args.port))
 
-    address = (args.hostname, int(args.port))
+        if args.udp:
+            connection = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+            try:
+                connection.bind(bind_addr)
+            except OSError:
+                print_error(f"Could not bind UDP listener on {bind_addr[0]}:{bind_addr[1]}")
+            print(Fore.GREEN + Style.BRIGHT + f"Listening (UDP) on {bind_addr[0]}:{bind_addr[1]}" + Style.RESET_ALL)
+        else:
+            server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            try:
+                server.bind(bind_addr)
+                server.listen(1)
+            except OSError:
+                print_error(f"Could not bind TCP listener on {bind_addr[0]}:{bind_addr[1]}")
+            print(Fore.GREEN + Style.BRIGHT + f"Listening (TCP) on {bind_addr[0]}:{bind_addr[1]}" + Style.RESET_ALL)
+            conn, addr = server.accept()
+            server.close()
+            connection = conn
+            peer = addr
+            print(Fore.GREEN + Style.BRIGHT + f"Connection established from {addr[0]}:{addr[1]}" + Style.RESET_ALL)
+    else:
+        if args.udp:
+            connection = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        else:
+            connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 
-    try:
-        connection.connect(address)
+        connection.settimeout(2)
+
+        address = (args.hostname, int(args.port))
 
-    except socket.error:
-        print_error("Could not establish connection to " + address[0] + ":" + str(address[1]))
+        try:
+            connection.connect(address)
+        except OSError:
+            print_error("Could not establish connection to " + address[0] + ":" + str(address[1]))
 
-    print(Fore.GREEN + Style.BRIGHT + "Connection established" + Style.RESET_ALL)
+        print(Fore.GREEN + Style.BRIGHT + "Connection established" + Style.RESET_ALL)
 
     try:
         connection.setblocking(0)
         stdin = ReadAsync(sys.stdin.readline)
 
         while True:
             try:
-                data = connection.recv(4096)
-                if not data:
-                    raise socket.error
-                print_ascii(data)
-                print_hex(to_hex(data))
-            except socket.error, e:
-                if e.errno != errno.EWOULDBLOCK:
+                if args.listen and args.udp:
+                    data, addr = connection.recvfrom(4096)
+                    if data:
+                        peer = addr
+                        print_ascii(data)
+                        print_hex(to_hex(data))
+                else:
+                    data = connection.recv(4096)
+                    if not data:
+                        raise OSError
+                    print_ascii(data)
+                    print_hex(to_hex(data))
+            except OSError as e:
+                if getattr(e, "errno", None) not in (errno.EWOULDBLOCK, errno.EAGAIN):
                     raise
             try:
-                connection.send(stdin.dequeue())
+                outbound = stdin.dequeue()
+                if outbound == "":
+                    # EOF on stdin
+                    raise KeyboardInterrupt
+
+                if args.send_hex:
+                    try:
+                        outbound_bytes = parse_hex_bytes(outbound.strip())
+                    except ValueError as e:
+                        print_error(f"Invalid hex input: {e}")
+                    if args.listen and args.udp:
+                        if peer is None:
+                            print_error("No UDP peer yet (send requires receiving at least one datagram)")
+                        connection.sendto(outbound_bytes, peer)
+                    else:
+                        connection.send(outbound_bytes)
+                else:
+                    if isinstance(outbound, str):
+                        outbound = outbound.encode("utf-8")
+                    if args.listen and args.udp:
+                        if peer is None:
+                            print_error("No UDP peer yet (send requires receiving at least one datagram)")
+                        connection.sendto(outbound, peer)
+                    else:
+                        connection.send(outbound)
             except Empty:
                 time.sleep(0.1)
+            except (BlockingIOError, InterruptedError):
+                time.sleep(0.01)
 
     except KeyboardInterrupt:
         connection.close()
         print_error("\nExiting...")
-    except socket.error:
+    except OSError:
         print_error("Connection closed")