Merge remote-tracking branch 'upstream/master'

5 files changed, 71 insertions, 62 deletions

diff --git a/pathod/language/base.py b/pathod/language/base.py
index 1369a3c7..25f3fd1a 100644
--- a/pathod/language/base.py
+++ b/pathod/language/base.py
@@ -261,7 +261,7 @@ class _Component(Token):
 
     """
         A value component of the primary specification of an message.
-        Components produce byte values desribe the bytes of the message.
+        Components produce byte values describing the bytes of the message.
     """
 
     def values(self, settings):  # pragma: no cover
@@ -272,9 +272,9 @@ class _Component(Token):
 
     def string(self, settings=None):
         """
-            A string representation of the object.
+            A bytestring representation of the object.
         """
-        return "".join(i[:] for i in self.values(settings or {}))
+        return b"".join(i[:] for i in self.values(settings or {}))
 
 
 class KeyValue(_Component):
@@ -391,7 +391,7 @@ class Integer(_Component):
                 "Integer value must be between %s and %s." % self.bounds,
                 0, 0
             )
-        self.value = str(value)
+        self.value = str(value).encode()
 
     @classmethod
     def expr(cls):
@@ -401,10 +401,10 @@ class Integer(_Component):
         return e.setParseAction(lambda x: cls(*x))
 
     def values(self, settings):
-        return self.value
+        return [self.value]
 
     def spec(self):
-        return "%s%s" % (self.preamble, self.value)
+        return "%s%s" % (self.preamble, self.value.decode())
 
     def freeze(self, settings_):
         return self
@@ -555,7 +555,7 @@ class NestedMessage(Token):
         try:
             self.parsed = self.nest_type(
                 self.nest_type.expr().parseString(
-                    value.val,
+                    value.val.decode(),
                     parseAll=True
                 )
             )
@@ -578,4 +578,4 @@ class NestedMessage(Token):
 
     def freeze(self, settings):
         f = self.parsed.freeze(settings).spec()
-        return self.__class__(TokValueLiteral(strutils.bytes_to_escaped_str(f)))
+        return self.__class__(TokValueLiteral(strutils.bytes_to_escaped_str(f.encode())))
diff --git a/pathod/language/generators.py b/pathod/language/generators.py
index 9fff3082..68705941 100644
--- a/pathod/language/generators.py
+++ b/pathod/language/generators.py
@@ -3,6 +3,7 @@ import random
 import mmap
 
 import six
+import sys
 
 DATATYPES = dict(
     ascii_letters=string.ascii_letters.encode(),
@@ -68,7 +69,7 @@ class RandomGenerator(object):
     def __getitem__(self, x):
         chars = DATATYPES[self.dtype]
         if isinstance(x, slice):
-            return b"".join(rand_byte(chars) for _ in range(*x.indices(self.length)))
+            return b"".join(rand_byte(chars) for _ in range(*x.indices(min(self.length, sys.maxsize))))
         return rand_byte(chars)
 
     def __repr__(self):
diff --git a/pathod/language/http.py b/pathod/language/http.py
index b2308d5e..4cc7db5f 100644
--- a/pathod/language/http.py
+++ b/pathod/language/http.py
@@ -57,7 +57,7 @@ class _HeaderMixin(object):
     unique_name = None
 
     def format_header(self, key, value):
-        return [key, ": ", value, "\r\n"]
+        return [key, b": ", value, b"\r\n"]
 
     def values(self, settings):
         return self.format_header(
@@ -88,7 +88,7 @@ class ShortcutUserAgent(_HeaderMixin, base.OptionsOrValue):
     def values(self, settings):
         value = self.value.val
         if self.option_used:
-            value = user_agents.get_by_shortcut(value.lower())[2]
+            value = user_agents.get_by_shortcut(value.lower().decode())[2].encode()
 
         return self.format_header(
             self.key.get_generator(settings),
@@ -109,7 +109,7 @@ def get_header(val, headers):
 
 
 class _HTTPMessage(message.Message):
-    version = "HTTP/1.1"
+    version = b"HTTP/1.1"
 
     @property
     def actions(self):
@@ -133,10 +133,10 @@ class _HTTPMessage(message.Message):
 
     def values(self, settings):
         vals = self.preamble(settings)
-        vals.append("\r\n")
+        vals.append(b"\r\n")
         for h in self.headers:
             vals.extend(h.values(settings))
-        vals.append("\r\n")
+        vals.append(b"\r\n")
         if self.body:
             vals.extend(self.body.values(settings))
         return vals
@@ -171,18 +171,18 @@ class Response(_HTTPMessage):
         return self.tok(Reason)
 
     def preamble(self, settings):
-        l = [self.version, " "]
+        l = [self.version, b" "]
         l.extend(self.status_code.values(settings))
         status_code = int(self.status_code.value)
-        l.append(" ")
+        l.append(b" ")
         if self.reason:
             l.extend(self.reason.values(settings))
         else:
             l.append(
                 status_codes.RESPONSES.get(
                     status_code,
-                    "Unknown code"
-                )
+                    b"Unknown code"
+                ).encode()
             )
         return l
 
@@ -205,8 +205,8 @@ class Response(_HTTPMessage):
                 if not get_header(i[0], self.headers):
                     tokens.append(
                         Header(
-                            base.TokValueLiteral(i[0]),
-                            base.TokValueLiteral(i[1]))
+                            base.TokValueLiteral(i[0].decode()),
+                            base.TokValueLiteral(i[1].decode()))
                     )
         if not self.raw:
             if not get_header("Content-Length", self.headers):
@@ -294,11 +294,11 @@ class Request(_HTTPMessage):
 
     def preamble(self, settings):
         v = self.method.values(settings)
-        v.append(" ")
+        v.append(b" ")
         v.extend(self.path.values(settings))
         if self.nested_response:
             v.append(self.nested_response.parsed.spec())
-        v.append(" ")
+        v.append(b" ")
         v.append(self.version)
         return v
 
@@ -314,8 +314,8 @@ class Request(_HTTPMessage):
                 if not get_header(i[0], self.headers):
                     tokens.append(
                         Header(
-                            base.TokValueLiteral(i[0]),
-                            base.TokValueLiteral(i[1])
+                            base.TokValueLiteral(i[0].decode()),
+                            base.TokValueLiteral(i[1].decode())
                         )
                     )
         if not self.raw:
diff --git a/pathod/language/message.py b/pathod/language/message.py
index 33124856..fea4f4de 100644
--- a/pathod/language/message.py
+++ b/pathod/language/message.py
@@ -1,5 +1,6 @@
 import abc
 from . import actions, exceptions
+from netlib import strutils
 
 LOG_TRUNCATE = 1024
 
@@ -49,7 +50,7 @@ class Message(object):
 
     def preview_safe(self):
         """
-            Return a copy of this message that issafe for previews.
+            Return a copy of this message that is safe for previews.
         """
         tokens = [i for i in self.tokens if not isinstance(i, actions.PauseAt)]
         return self.__class__(tokens)
@@ -80,10 +81,10 @@ class Message(object):
             # We truncate at 1k.
             if hasattr(v, "values"):
                 v = [x[:LOG_TRUNCATE] for x in v.values(settings)]
-                v = "".join(v).encode("string_escape")
+                v = strutils.bytes_to_escaped_str(b"".join(v))
             elif hasattr(v, "__len__"):
                 v = v[:LOG_TRUNCATE]
-                v = v.encode("string_escape")
+                v = strutils.bytes_to_escaped_str(v)
             ret[i] = v
         ret["spec"] = self.spec()
         return ret
diff --git a/pathod/pathoc.py b/pathod/pathoc.py
index 5cfb4591..3eb91637 100644
--- a/pathod/pathoc.py
+++ b/pathod/pathoc.py
@@ -247,7 +247,10 @@ class Pathoc(tcp.TCPClient):
 
     def socks_connect(self, connect_to):
         try:
-            client_greet = socks.ClientGreeting(socks.VERSION.SOCKS5, [socks.METHOD.NO_AUTHENTICATION_REQUIRED])
+            client_greet = socks.ClientGreeting(
+                socks.VERSION.SOCKS5,
+                [socks.METHOD.NO_AUTHENTICATION_REQUIRED]
+            )
             client_greet.to_file(self.wfile)
             self.wfile.flush()
 
@@ -286,44 +289,47 @@ class Pathoc(tcp.TCPClient):
         if self.use_http2 and not self.ssl:
             raise NotImplementedError("HTTP2 without SSL is not supported.")
 
-        ret = tcp.TCPClient.connect(self)
+        try:
+            ret = tcp.TCPClient.connect(self)
+            if connect_to:
+                self.http_connect(connect_to)
 
-        if connect_to:
-            self.http_connect(connect_to)
+            self.sslinfo = None
+            if self.ssl:
+                try:
+                    alpn_protos = [b'http/1.1']
+                    if self.use_http2:
+                        alpn_protos.append(b'h2')
+
+                    self.convert_to_ssl(
+                        sni=self.sni,
+                        cert=self.clientcert,
+                        method=self.ssl_version,
+                        options=self.ssl_options,
+                        cipher_list=self.ciphers,
+                        alpn_protos=alpn_protos
+                    )
+                except TlsException as v:
+                    raise PathocError(str(v))
 
-        self.sslinfo = None
-        if self.ssl:
-            try:
-                alpn_protos = [b'http/1.1']
-                if self.use_http2:
-                    alpn_protos.append(b'h2')
-
-                self.convert_to_ssl(
-                    sni=self.sni,
-                    cert=self.clientcert,
-                    method=self.ssl_version,
-                    options=self.ssl_options,
-                    cipher_list=self.ciphers,
-                    alpn_protos=alpn_protos
+                self.sslinfo = SSLInfo(
+                    self.connection.get_peer_cert_chain(),
+                    self.get_current_cipher(),
+                    self.get_alpn_proto_negotiated()
                 )
-            except TlsException as v:
-                raise PathocError(str(v))
-
-            self.sslinfo = SSLInfo(
-                self.connection.get_peer_cert_chain(),
-                self.get_current_cipher(),
-                self.get_alpn_proto_negotiated()
-            )
-            if showssl:
-                print(str(self.sslinfo), file=fp)
-
-            if self.use_http2:
-                self.protocol.check_alpn()
-                if not self.http2_skip_connection_preface:
-                    self.protocol.perform_client_connection_preface()
+                if showssl:
+                    print(str(self.sslinfo), file=fp)
 
-        if self.timeout:
-            self.settimeout(self.timeout)
+                if self.use_http2:
+                    self.protocol.check_alpn()
+                    if not self.http2_skip_connection_preface:
+                        self.protocol.perform_client_connection_preface()
+
+            if self.timeout:
+                self.settimeout(self.timeout)
+        except Exception:
+            self.close()
+            raise
         return ret
 
     def stop(self):
@@ -355,6 +361,7 @@ class Pathoc(tcp.TCPClient):
                         return
                 if frm is None:
                     self.ws_framereader.join()
+                    self.ws_framereader = None
                     return
                 yield frm