python3++

3 files changed, 22 insertions, 12 deletions

diff --git a/netlib/certutils.py b/netlib/certutils.py
index df793537..b3ddcbe4 100644
--- a/netlib/certutils.py
+++ b/netlib/certutils.py
@@ -3,7 +3,7 @@ import os
 import ssl
 import time
 import datetime
-import itertools
+from six.moves import filter
 import ipaddress
 
 import sys
@@ -396,12 +396,12 @@ class SSLCert(object):
     @property
     def notbefore(self):
         t = self.x509.get_notBefore()
-        return datetime.datetime.strptime(t, "%Y%m%d%H%M%SZ")
+        return datetime.datetime.strptime(t.decode("ascii"), "%Y%m%d%H%M%SZ")
 
     @property
     def notafter(self):
         t = self.x509.get_notAfter()
-        return datetime.datetime.strptime(t, "%Y%m%d%H%M%SZ")
+        return datetime.datetime.strptime(t.decode("ascii"), "%Y%m%d%H%M%SZ")
 
     @property
     def has_expired(self):
diff --git a/netlib/socks.py b/netlib/socks.py
index d38b88c8..51ad1c63 100644
--- a/netlib/socks.py
+++ b/netlib/socks.py
@@ -1,7 +1,7 @@
 from __future__ import (absolute_import, print_function, division)
-import socket
 import struct
 import array
+import ipaddress
 from . import tcp, utils
 
 
@@ -133,19 +133,23 @@ class Message(object):
     def from_file(cls, f):
         ver, msg, rsv, atyp = struct.unpack("!BBBB", f.safe_read(4))
         if rsv != 0x00:
-            raise SocksError(REP.GENERAL_SOCKS_SERVER_FAILURE,
-                             "Socks Request: Invalid reserved byte: %s" % rsv)
-
+            raise SocksError(
+                REP.GENERAL_SOCKS_SERVER_FAILURE,
+                "Socks Request: Invalid reserved byte: %s" % rsv
+            )
         if atyp == ATYP.IPV4_ADDRESS:
             # We use tnoa here as ntop is not commonly available on Windows.
-            host = socket.inet_ntoa(f.safe_read(4))
+            host = ipaddress.IPv4Address(f.safe_read(4)).compressed
             use_ipv6 = False
         elif atyp == ATYP.IPV6_ADDRESS:
-            host = socket.inet_ntop(socket.AF_INET6, f.safe_read(16))
+            host = ipaddress.IPv6Address(f.safe_read(16)).compressed
             use_ipv6 = True
         elif atyp == ATYP.DOMAINNAME:
             length, = struct.unpack("!B", f.safe_read(1))
             host = f.safe_read(length)
+            if not utils.is_valid_host(host):
+                raise SocksError(REP.GENERAL_SOCKS_SERVER_FAILURE, "Invalid hostname: %s" % host)
+            host = host.decode("idna")
             use_ipv6 = False
         else:
             raise SocksError(REP.ADDRESS_TYPE_NOT_SUPPORTED,
@@ -158,12 +162,12 @@ class Message(object):
     def to_file(self, f):
         f.write(struct.pack("!BBBB", self.ver, self.msg, 0x00, self.atyp))
         if self.atyp == ATYP.IPV4_ADDRESS:
-            f.write(socket.inet_aton(self.addr.host))
+            f.write(ipaddress.IPv4Address(self.addr.host).packed)
         elif self.atyp == ATYP.IPV6_ADDRESS:
-            f.write(socket.inet_pton(socket.AF_INET6, self.addr.host))
+            f.write(ipaddress.IPv6Address(self.addr.host).packed)
         elif self.atyp == ATYP.DOMAINNAME:
             f.write(struct.pack("!B", len(self.addr.host)))
-            f.write(self.addr.host)
+            f.write(self.addr.host.encode("idna"))
         else:
             raise SocksError(
                 REP.ADDRESS_TYPE_NOT_SUPPORTED,
diff --git a/netlib/utils.py b/netlib/utils.py
index 6fed44b6..799b0d42 100644
--- a/netlib/utils.py
+++ b/netlib/utils.py
@@ -141,6 +141,12 @@ _label_valid = re.compile(b"(?!-)[A-Z\d-]{1,63}(?<!-)$", re.IGNORECASE)
 
 
 def is_valid_host(host):
+    """
+    Checks if a hostname is valid.
+
+    Args:
+      host (bytes): The hostname
+    """
     try:
         host.decode("idna")
     except ValueError: