From 30681a701869b0d8431354184f1c5106eacdf5c2 Mon Sep 17 00:00:00 2001
From: Algervivien <darthabel@kadmony>
Date: Mon, 27 Jun 2011 18:49:01 +0200
Subject: [PATCH] Use of socket.getaddrinfo() to replace ipv6 attributes.
 Corrected error with ssl-only option

---
 websocket.py | 21 ++++++++++-----------
 websockify   | 24 +++++++-----------------
 2 files changed, 17 insertions(+), 28 deletions(-)

diff --git a/websocket.py b/websocket.py
index cd94fde..78fc586 100755
--- a/websocket.py
+++ b/websocket.py
@@ -95,7 +95,6 @@ Sec-WebSocket-Accept: %s\r
         self.verbose        = verbose
         self.listen_host    = listen_host
         self.listen_port    = listen_port
-        self.source_is_ipv6 = source_is_ipv6
         self.ssl_only       = ssl_only
         self.daemon         = daemon
         self.handler_id     = 1
@@ -114,7 +113,7 @@ Sec-WebSocket-Accept: %s\r
             os.chdir(self.web)
 
         # Sanity checks
-        if ssl and self.ssl_only:
+        if not ssl and self.ssl_only:
             raise Exception("No 'ssl' module and SSL-only specified")
         if self.daemon and not resource:
             raise Exception("Module 'resource' required to daemonize")
@@ -512,6 +511,7 @@ Sec-WebSocket-Accept: %s\r
             if not os.path.exists(self.cert):
                 raise self.EClose("SSL connection but '%s' not found"
                                   % self.cert)
+            retsock = None
             try:
                 retsock = ssl.wrap_socket(
                         sock,
@@ -687,15 +687,14 @@ Sec-WebSocket-Accept: %s\r
         is a WebSockets client then call new_client() method (which must
         be overridden) for each new client connection.
         """
-        
-        if self.source_is_ipv6:
-          lsock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
-        else:
-          lsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        lsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        lsock.bind((self.listen_host, self.listen_port))
-        lsock.listen(100)
-
+        for res in socket.getaddrinfo(self.listen_host, self.listen_port, 0, socket.SOCK_STREAM, 6):
+            af, socktype, proto, canonname, sa = res
+            lsock = socket.socket(af, socktype)
+            lsock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            lsock.bind((self.listen_host, self.listen_port))
+            lsock.listen(100)
+            break
+            
         if self.daemon:
             self.daemonize(keepfd=lsock.fileno(), chdir=self.web)
 
diff --git a/websockify b/websockify
index 126f7ec..8d1fe0f 100755
--- a/websockify
+++ b/websockify
@@ -41,7 +41,6 @@ Traffic Legend:
         # Save off proxy specific options
         self.target_host    = kwargs.pop('target_host')
         self.target_port    = kwargs.pop('target_port')
-        self.target_is_ipv6 = kwargs.pop('target_is_ipv6')
         self.wrap_cmd       = kwargs.pop('wrap_cmd')
         self.wrap_mode      = kwargs.pop('wrap_mode')
         # Last 3 timestamps command was run
@@ -151,12 +150,12 @@ Traffic Legend:
         # Connect to the target
         self.msg("connecting to: %s:%s" % (
                  self.target_host, self.target_port))
-        if self.target_is_ipv6:
-            tsock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
-        else:
-            tsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        tsock.connect((self.target_host, self.target_port))
-
+        for res in socket.getaddrinfo(self.target_host, self.target_port, 0, socket.SOCK_STREAM, 6):
+            af, socktype, proto, canonname, sa = res
+            tsock = socket.socket(af, socktype)
+            tsock.connect((self.target_host, self.target_port))
+            break
+                        
         if self.verbose and not self.daemon:
             print(self.traffic_legend)
 
@@ -283,11 +282,7 @@ if __name__ == '__main__':
         opts.listen_host, sep, opts.listen_port = args[0].rpartition(':')
     else:
         opts.listen_host, opts.listen_port = '', args[0]
-    if opts.listen_host.count(':') > 0:
-        opts.source_is_ipv6 = True
-    else:
-        opts.source_is_ipv6 = False
-
+    
     try:    opts.listen_port = int(opts.listen_port)
     except: parser.error("Error parsing listen port")
     
@@ -297,11 +292,6 @@ if __name__ == '__main__':
     else:
         if args[1].count(':') > 0:
             opts.target_host, sep, opts.target_port = args[1].rpartition(':')
-            if opts.target_host.count(':') > 0:
-                opts.target_is_ipv6 = True
-            else:
-                opts.target_is_ipv6 = False
-
         else:
             parser.error("Error parsing target")
         try:    opts.target_port = int(opts.target_port)