diff --git a/squinnondation/squinnondation.py b/squinnondation/squinnondation.py index 6cbd118..073444d 100644 --- a/squinnondation/squinnondation.py +++ b/squinnondation/squinnondation.py @@ -4,6 +4,7 @@ import socket from argparse import ArgumentParser from enum import Enum +from ipaddress import IPv6Address from typing import Any, Optional, Tuple @@ -170,19 +171,19 @@ class HelloTLV(TLV): class NeighbourTLV(TLV): type: int = 3 length: int - ip_address: int + ip_address: IPv6Address port: int def unmarshal(self, raw_data: bytes) -> None: self.type = raw_data[0] self.length = raw_data[1] - self.ip_address = raw_data[2:18] + self.ip_address = IPv6Address(raw_data[2:18]) self.port = int.from_bytes(raw_data[18:20], "big") def marshal(self) -> bytes: return self.type.to_bytes(1, "big") + \ self.length.to_bytes(1, "big") + \ - self.ip_address.to_bytes(16, "big") + \ + self.ip_address.packed + \ self.port.to_bytes(2, "big") @@ -328,7 +329,9 @@ class Hazelnut: """ def __init__(self, nickname: str = "anonymous", address: str = "localhost", port: int = 2500): self.nickname = nickname - self.address = address + # Resolve DNS as an IPv6 + address = socket.getaddrinfo(address, None, socket.AF_INET6)[0][4][0] + self.address = IPv6Address(address) self.port = port @@ -339,7 +342,7 @@ class Squirrel(Hazelnut): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.socket = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) - self.socket.bind((self.address, self.port)) + self.socket.bind((str(self.address), self.port)) print(f"Listening on {self.address}:{self.port}") def send_packet(self, client: Hazelnut, pkt: Packet) -> int: @@ -352,7 +355,7 @@ class Squirrel(Hazelnut): """ Send a raw packet to a client. """ - return self.socket.sendto(data, (client.address, client.port)) + return self.socket.sendto(data, (str(client.address), client.port)) def receive_packet(self) -> Tuple[Packet, Any]: """