Resolve DNS as IPv6 addresses before connecting to the socket

Signed-off-by: Yohann D'ANELLO <ynerant@crans.org>
This commit is contained in:
Yohann D'ANELLO 2020-12-21 16:03:58 +01:00
parent cf4833573a
commit 79b8a44461
Signed by: ynerant
GPG Key ID: 3A75C55819C8CF85
1 changed files with 9 additions and 6 deletions

View File

@ -4,6 +4,7 @@
import socket import socket
from argparse import ArgumentParser from argparse import ArgumentParser
from enum import Enum from enum import Enum
from ipaddress import IPv6Address
from typing import Any, Optional, Tuple from typing import Any, Optional, Tuple
@ -170,19 +171,19 @@ class HelloTLV(TLV):
class NeighbourTLV(TLV): class NeighbourTLV(TLV):
type: int = 3 type: int = 3
length: int length: int
ip_address: int ip_address: IPv6Address
port: int port: int
def unmarshal(self, raw_data: bytes) -> None: def unmarshal(self, raw_data: bytes) -> None:
self.type = raw_data[0] self.type = raw_data[0]
self.length = raw_data[1] 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") self.port = int.from_bytes(raw_data[18:20], "big")
def marshal(self) -> bytes: def marshal(self) -> bytes:
return self.type.to_bytes(1, "big") + \ return self.type.to_bytes(1, "big") + \
self.length.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") self.port.to_bytes(2, "big")
@ -328,7 +329,9 @@ class Hazelnut:
""" """
def __init__(self, nickname: str = "anonymous", address: str = "localhost", port: int = 2500): def __init__(self, nickname: str = "anonymous", address: str = "localhost", port: int = 2500):
self.nickname = nickname 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 self.port = port
@ -339,7 +342,7 @@ class Squirrel(Hazelnut):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
self.socket = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) 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}") print(f"Listening on {self.address}:{self.port}")
def send_packet(self, client: Hazelnut, pkt: Packet) -> int: def send_packet(self, client: Hazelnut, pkt: Packet) -> int:
@ -352,7 +355,7 @@ class Squirrel(Hazelnut):
""" """
Send a raw packet to a client. 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]: def receive_packet(self) -> Tuple[Packet, Any]:
""" """