support for IPv6 (client side only)

tftpy/TftpContexts.py

@@ -83,19 +83,20 @@ def __init__(self, host, port, timeout, localip = ""):
         self.fileobj = None
         self.options = None
         self.packethook = None
-        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        # Note, setting the host will also set self.address, as it's a property.
+        self.port = port
+        self.host = host
+        # The port associated with the TID
+        self.tidport = None
+        # Socket
+        self.sock = socket.socket(self.family, socket.SOCK_DGRAM)
         if localip != "":
             self.sock.bind((localip, 0))
         self.sock.settimeout(timeout)
         self.timeout = timeout
         self.state = None
         self.next_block = 0
         self.factory = TftpPacketFactory()
-        # Note, setting the host will also set self.address, as it's a property.
-        self.host = host
-        self.port = port
-        # The port associated with the TID
-        self.tidport = None
         # Metrics
         self.metrics = TftpMetrics()
         # Fluag when the transfer is pending completion.
@@ -144,10 +145,17 @@ def gethost(self):
         return self.__host
 
     def sethost(self, host):
-        """Setter method that also sets the address property as a result
-        of the host that is set."""
-        self.__host = host
-        self.address = socket.gethostbyname(host)
+        """Setter method that also sets the address and family properties
+        as a result of the host that is set."""
+        try:
+            self.__host = host
+            self.address = socket.gethostbyname(host)
+            self.family = socket.AF_INET
+        except socket.gaierror:
+            (family, type, proto, canonname, sockaddr) = socket.getaddrinfo(host, self.port)[0]
+            self.__host = canonname if canonname else sockaddr[0]
+            self.address = sockaddr[0]
+            self.family = family
 
     host = property(gethost, sethost)
 
@@ -166,7 +174,10 @@ def cycle(self):
         """Here we wait for a response from the server after sending it
         something, and dispatch appropriate action to that response."""
         try:
-            (buffer, (raddress, rport)) = self.sock.recvfrom(MAX_BLKSIZE)
+            (buffer, address) = self.sock.recvfrom(MAX_BLKSIZE)
+            # IPv4 and IPv6 sockets return a different number of flags
+            raddress = address[0]
+            rport = address[1]
         except socket.timeout:
             log.warning("Timeout waiting for traffic, retrying...")
             raise TftpTimeout("Timed-out waiting for traffic")