From 41140b8a36495c1ea1769a777c3a84b3919f95e5 Mon Sep 17 00:00:00 2001
From: Devin Neal <devneal17@gmail.com>
Date: Mon, 13 Nov 2017 15:17:00 -0500
Subject: [PATCH 1/8] add server module

---
 pwn/toplevel.py        |   1 +
 pwnlib/tubes/server.py | 149 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 150 insertions(+)
 create mode 100644 pwnlib/tubes/server.py

diff --git a/pwn/toplevel.py b/pwn/toplevel.py
index 14ed3a82a..9dac9897c 100644
--- a/pwn/toplevel.py
+++ b/pwn/toplevel.py
@@ -45,6 +45,7 @@
 from pwnlib.tubes.serialtube import serialtube
 from pwnlib.tubes.ssh import ssh
 from pwnlib.tubes.tube import tube
+from pwnlib.tubes.server import server
 from pwnlib.ui import *
 from pwnlib.util import crc
 from pwnlib.util import iters
diff --git a/pwnlib/tubes/server.py b/pwnlib/tubes/server.py
new file mode 100644
index 000000000..6a4f6ef14
--- /dev/null
+++ b/pwnlib/tubes/server.py
@@ -0,0 +1,149 @@
+from __future__ import absolute_import
+
+import errno
+import socket
+
+from pwnlib.context import context
+from pwnlib.log import getLogger
+from pwnlib.timeout import Timeout
+from pwnlib.tubes.sock import sock
+from pwnlib.tubes.remote import remote
+
+log = getLogger(__name__)
+
+class server(sock):
+    r"""Creates an TCP or UDP-server to listen for connections. It supports
+    both IPv4 and IPv6.
+
+    The returned object supports all the methods from
+    The callback should take a :class:`pwnlib.tubes.remote` as
+    it's only argument.
+
+    Arguments:
+        port(int): The port to connect to.
+            Defaults to a port auto-selected by the operating system.
+        bindaddr(str): The address to bind to.
+            Defaults to ``0.0.0.0`` / `::`.
+        fam: The string "any", "ipv4" or "ipv6" or an integer to pass to :func:`socket.getaddrinfo`.
+        typ: The string "tcp" or "udp" or an integer to pass to :func:`socket.getaddrinfo`.
+
+    Examples:
+
+        >>> s = server(8888)
+        >>> r1 = remote('localhost', s.lport)
+        >>> r2 = remote('localhost', s.lport)
+        >>> r1.sendline("Hello")
+        >>> r2.sendline("Hi")
+        >>> r2.recvline()
+        'Hi\n'
+        >>> r1.recvline()
+        'Hello\n'
+    """
+
+    #: Local port
+    lport = 0
+
+    #: Local host
+    lhost = None
+
+    #: Socket type (e.g. socket.SOCK_STREAM)
+    type = None
+
+    #: Socket family
+    family = None
+
+    #: Socket protocol
+    protocol = None
+
+    #: Canonical name of the listening interface
+    canonname = None
+
+    #: Sockaddr structure that is being listened on
+    sockaddr = None
+
+    _accepter = None
+
+    def __init__(self, port=0, bindaddr = "0.0.0.0",
+                 fam = "any", typ = "tcp", callback = None, *args, **kwargs):
+        super(server, self).__init__(*args, **kwargs)
+
+        port = int(port)
+        fam  = {socket.AF_INET: 'ipv4',
+                socket.AF_INET6: 'ipv6'}.get(fam, fam)
+
+        fam = self._get_family(fam)
+        typ = self._get_type(typ)
+
+        if fam == socket.AF_INET6 and bindaddr == '0.0.0.0':
+            bindaddr = '::'
+
+        def echo(remote):
+            while True:
+                s = remote.readline()
+                remote.send(s)
+
+        if not callback:
+            callback = echo
+
+        h = self.waitfor('Trying to bind to %s on port %d' % (bindaddr, port))
+
+        for res in socket.getaddrinfo(bindaddr, port, fam, typ, 0, socket.AI_PASSIVE):
+            self.family, self.type, self.proto, self.canonname, self.sockaddr = res
+
+            if self.type not in [socket.SOCK_STREAM, socket.SOCK_DGRAM]:
+                continue
+
+            h.status("Trying %s" % self.sockaddr[0])
+            listen_sock = socket.socket(self.family, self.type, self.proto)
+            listen_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            listen_sock.bind(self.sockaddr)
+            self.lhost, self.lport = listen_sock.getsockname()[:2]
+            if self.type == socket.SOCK_STREAM:
+                listen_sock.listen(1)
+            break
+        else:
+            h.failure()
+            self.error("Could not bind to %s on port %d" % (bindaddr, port))
+
+        h.success()
+
+        def accepter():
+            while True:
+                h = self.waitfor('Waiting for connections on %s:%s' % (self.lhost, self.lport))
+                while True:
+                    try:
+                        if self.type == socket.SOCK_STREAM:
+                            sock, rhost = listen_sock.accept()
+                            rhost, rport = rhost[:2]
+                        else:
+                            data, rhost = listen_sock.recvfrom(4096)
+                            listen_sock.connect(rhost)
+                            sock = listen_sock
+                            sock.unrecv(data)
+                        sock.settimeout(self.timeout)
+                        break
+                    except socket.error as e:
+                        if e.errno == errno.EINTR:
+                            continue
+                        h.failure()
+                        self.exception("Socket failure while waiting for connection")
+                        sock = None
+                        return
+
+                r = remote(rhost, rport, sock = sock)
+                t = context.Thread(target = callback, args = (r,))
+                t.daemon = False
+                t.start()
+                h.success('Got connection from %s on port %d' % (rhost, rport))
+
+        self._accepter = context.Thread(target = accepter)
+        self._accepter.daemon = False
+        self._accepter.start()
+
+    def close(self):
+        self._accepter.close()
+        # since `close` is scheduled to run on exit we must check that we got
+        # a connection or the program will hang in the `join` call above
+        if self._accepter and self._accepter.is_alive():
+            return
+        super(server, self).close()

From 2deb9fe5271171f1115409819448a784aca5a129 Mon Sep 17 00:00:00 2001
From: Devin Neal <devneal17@gmail.com>
Date: Mon, 13 Nov 2017 15:33:37 -0500
Subject: [PATCH 2/8] modify example

---
 pwnlib/tubes/server.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/pwnlib/tubes/server.py b/pwnlib/tubes/server.py
index 6a4f6ef14..22414fe7e 100644
--- a/pwnlib/tubes/server.py
+++ b/pwnlib/tubes/server.py
@@ -15,9 +15,8 @@ class server(sock):
     r"""Creates an TCP or UDP-server to listen for connections. It supports
     both IPv4 and IPv6.
 
-    The returned object supports all the methods from
-    The callback should take a :class:`pwnlib.tubes.remote` as
-    it's only argument.
+    The callback function should take a :class:`pwnlib.tubes.remote` as
+    its only argument.
 
     Arguments:
         port(int): The port to connect to.
@@ -26,6 +25,7 @@ class server(sock):
             Defaults to ``0.0.0.0`` / `::`.
         fam: The string "any", "ipv4" or "ipv6" or an integer to pass to :func:`socket.getaddrinfo`.
         typ: The string "tcp" or "udp" or an integer to pass to :func:`socket.getaddrinfo`.
+        callback: The function to run with the new connection.
 
     Examples:
 
@@ -38,6 +38,15 @@ class server(sock):
         'Hi\n'
         >>> r1.recvline()
         'Hello\n'
+        >>> def callback(conn):
+        ...     s = conn.recvline()
+        ...     conn.send(s[::-1])
+        ... 
+        >>> t = server(8889, callback=callback)
+        >>> r3 = remote('localhost', t.lport)
+        >>> r3.sendline('callback')
+        >>> r3.recv()
+        '\nkcabllac'
     """
 
     #: Local port

From 5286d7c7204bca39edaa32cc21e45aab5e776831 Mon Sep 17 00:00:00 2001
From: Devin Neal <devneal17@gmail.com>
Date: Mon, 13 Nov 2017 16:46:55 -0500
Subject: [PATCH 3/8] fix travis errors

---
 pwn/toplevel.py        | 1 -
 pwnlib/tubes/server.py | 3 +--
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/pwn/toplevel.py b/pwn/toplevel.py
index 9dac9897c..14ed3a82a 100644
--- a/pwn/toplevel.py
+++ b/pwn/toplevel.py
@@ -45,7 +45,6 @@
 from pwnlib.tubes.serialtube import serialtube
 from pwnlib.tubes.ssh import ssh
 from pwnlib.tubes.tube import tube
-from pwnlib.tubes.server import server
 from pwnlib.ui import *
 from pwnlib.util import crc
 from pwnlib.util import iters
diff --git a/pwnlib/tubes/server.py b/pwnlib/tubes/server.py
index 22414fe7e..da7ea04e9 100644
--- a/pwnlib/tubes/server.py
+++ b/pwnlib/tubes/server.py
@@ -5,7 +5,6 @@
 
 from pwnlib.context import context
 from pwnlib.log import getLogger
-from pwnlib.timeout import Timeout
 from pwnlib.tubes.sock import sock
 from pwnlib.tubes.remote import remote
 
@@ -41,7 +40,7 @@ class server(sock):
         >>> def callback(conn):
         ...     s = conn.recvline()
         ...     conn.send(s[::-1])
-        ... 
+        ...
         >>> t = server(8889, callback=callback)
         >>> r3 = remote('localhost', t.lport)
         >>> r3.sendline('callback')

From 629a2b0bf30cf55cb2c74dbb48a03f30530f1e80 Mon Sep 17 00:00:00 2001
From: Devin Neal <devneal17@gmail.com>
Date: Tue, 14 Nov 2017 02:31:38 -0500
Subject: [PATCH 4/8] add queuing and blocking, fix requests

---
 pwnlib/tubes/server.py | 75 ++++++++++++++++++++++--------------------
 1 file changed, 40 insertions(+), 35 deletions(-)

diff --git a/pwnlib/tubes/server.py b/pwnlib/tubes/server.py
index da7ea04e9..b9c4d9bed 100644
--- a/pwnlib/tubes/server.py
+++ b/pwnlib/tubes/server.py
@@ -2,6 +2,7 @@
 
 import errno
 import socket
+import threading
 
 from pwnlib.context import context
 from pwnlib.log import getLogger
@@ -14,9 +15,6 @@ class server(sock):
     r"""Creates an TCP or UDP-server to listen for connections. It supports
     both IPv4 and IPv6.
 
-    The callback function should take a :class:`pwnlib.tubes.remote` as
-    its only argument.
-
     Arguments:
         port(int): The port to connect to.
             Defaults to a port auto-selected by the operating system.
@@ -24,27 +22,24 @@ class server(sock):
             Defaults to ``0.0.0.0`` / `::`.
         fam: The string "any", "ipv4" or "ipv6" or an integer to pass to :func:`socket.getaddrinfo`.
         typ: The string "tcp" or "udp" or an integer to pass to :func:`socket.getaddrinfo`.
-        callback: The function to run with the new connection.
+        callback: A function to be started on incoming connections. It should take a :class:`pwnlib.tubes.remote` as its only argument.
 
     Examples:
 
         >>> s = server(8888)
-        >>> r1 = remote('localhost', s.lport)
-        >>> r2 = remote('localhost', s.lport)
-        >>> r1.sendline("Hello")
-        >>> r2.sendline("Hi")
-        >>> r2.recvline()
-        'Hi\n'
-        >>> r1.recvline()
+        >>> client_conn = remote('localhost', s.lport)
+        >>> server_conn = s.next_connection()
+        >>> client_conn.sendline('Hello')
+        >>> server_conn.recvline()
         'Hello\n'
-        >>> def callback(conn):
-        ...     s = conn.recvline()
-        ...     conn.send(s[::-1])
+        >>> def cb(r):
+        ...     client_input = r.readline()
+        ...     r.send(client_input[::-1])
         ...
-        >>> t = server(8889, callback=callback)
-        >>> r3 = remote('localhost', t.lport)
-        >>> r3.sendline('callback')
-        >>> r3.recv()
+        >>> t = server(8889, callback=cb)
+        >>> client_conn = remote('localhost', t.lport)
+        >>> client_conn.sendline('callback')
+        >>> client_conn.recv()
         '\nkcabllac'
     """
 
@@ -71,8 +66,8 @@ class server(sock):
 
     _accepter = None
 
-    def __init__(self, port=0, bindaddr = "0.0.0.0",
-                 fam = "any", typ = "tcp", callback = None, *args, **kwargs):
+    def __init__(self, port=0, bindaddr = "0.0.0.0", fam = "any", typ = "tcp",
+                 callback = None, blocking = False, *args, **kwargs):
         super(server, self).__init__(*args, **kwargs)
 
         port = int(port)
@@ -85,14 +80,6 @@ def __init__(self, port=0, bindaddr = "0.0.0.0",
         if fam == socket.AF_INET6 and bindaddr == '0.0.0.0':
             bindaddr = '::'
 
-        def echo(remote):
-            while True:
-                s = remote.readline()
-                remote.send(s)
-
-        if not callback:
-            callback = echo
-
         h = self.waitfor('Trying to bind to %s on port %d' % (bindaddr, port))
 
         for res in socket.getaddrinfo(bindaddr, port, fam, typ, 0, socket.AI_PASSIVE):
@@ -115,6 +102,9 @@ def echo(remote):
 
         h.success()
 
+        self.sock = listen_sock
+        self.connections_waiting = threading.Event()
+        self.connections = []
         def accepter():
             while True:
                 h = self.waitfor('Waiting for connections on %s:%s' % (self.lhost, self.lport))
@@ -122,7 +112,6 @@ def accepter():
                     try:
                         if self.type == socket.SOCK_STREAM:
                             sock, rhost = listen_sock.accept()
-                            rhost, rport = rhost[:2]
                         else:
                             data, rhost = listen_sock.recvfrom(4096)
                             listen_sock.connect(rhost)
@@ -138,18 +127,34 @@ def accepter():
                         sock = None
                         return
 
-                r = remote(rhost, rport, sock = sock)
-                t = context.Thread(target = callback, args = (r,))
-                t.daemon = False
-                t.start()
-                h.success('Got connection from %s on port %d' % (rhost, rport))
+                self.rhost, self.rport = rhost[:2]
+                r = remote(self.rhost, self.rport, sock = sock)
+                h.success('Got connection from %s on port %d' % (self.rhost, self.rport))
+                if callback:
+                    if not blocking:
+                        t = context.Thread(target = callback, args = (r,))
+                        t.daemon = True
+                        t.start()
+                    else:
+                        callback(r)
+                else:
+                    self.connections.append(r)
+                    if not self.connections_waiting.is_set():
+                        self.connections_waiting.set()
 
         self._accepter = context.Thread(target = accepter)
         self._accepter.daemon = False
         self._accepter.start()
 
+    def next_connection(self):
+        if not self.connections_waiting.is_set():
+            self.connections_waiting.wait()
+        conn = self.connections.pop(0)
+        if not self.connections:
+            self.connections_waiting.clear()
+        return conn
+
     def close(self):
-        self._accepter.close()
         # since `close` is scheduled to run on exit we must check that we got
         # a connection or the program will hang in the `join` call above
         if self._accepter and self._accepter.is_alive():

From 187f2b709b87bc7bb9b6363d19e8adb04749768a Mon Sep 17 00:00:00 2001
From: Devin Neal <devneal17@gmail.com>
Date: Fri, 8 Dec 2017 17:24:24 -0500
Subject: [PATCH 5/8] add server.rst, change accepter to daemon

---
 docs/source/tubes/server.rst | 9 +++++++++
 pwnlib/tubes/server.py       | 2 +-
 2 files changed, 10 insertions(+), 1 deletion(-)
 create mode 100644 docs/source/tubes/server.rst

diff --git a/docs/source/tubes/server.rst b/docs/source/tubes/server.rst
new file mode 100644
index 000000000..720095f80
--- /dev/null
+++ b/docs/source/tubes/server.rst
@@ -0,0 +1,9 @@
+.. testsetup:: *
+
+   from pwnlib.tubes.server import server
+
+:mod:`pwnlib.tubes.server` --- Server
+==================================================
+
+.. automodule:: pwnlib.tubes.server
+   :members:
diff --git a/pwnlib/tubes/server.py b/pwnlib/tubes/server.py
index b9c4d9bed..e542c27ae 100644
--- a/pwnlib/tubes/server.py
+++ b/pwnlib/tubes/server.py
@@ -143,7 +143,7 @@ def accepter():
                         self.connections_waiting.set()
 
         self._accepter = context.Thread(target = accepter)
-        self._accepter.daemon = False
+        self._accepter.daemon = True
         self._accepter.start()
 
     def next_connection(self):

From 70948c06fdd520c631bcda5608124a12322cd75f Mon Sep 17 00:00:00 2001
From: Devin Neal <devneal17@gmail.com>
Date: Fri, 8 Dec 2017 18:30:55 -0500
Subject: [PATCH 6/8] fix import in server.rst

---
 docs/source/tubes/server.rst | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/source/tubes/server.rst b/docs/source/tubes/server.rst
index 720095f80..d6c6212d8 100644
--- a/docs/source/tubes/server.rst
+++ b/docs/source/tubes/server.rst
@@ -1,5 +1,6 @@
 .. testsetup:: *
 
+   from pwn import *
    from pwnlib.tubes.server import server
 
 :mod:`pwnlib.tubes.server` --- Server

From ae778b2a74c745b7bec3cab83006438f16239b40 Mon Sep 17 00:00:00 2001
From: Devin Neal <devneal17@gmail.com>
Date: Mon, 11 Dec 2017 14:45:47 -0500
Subject: [PATCH 7/8] added server entry in sockets.rst

---
 docs/source/tubes/server.rst  | 10 ----------
 docs/source/tubes/sockets.rst |  6 ++++++
 2 files changed, 6 insertions(+), 10 deletions(-)
 delete mode 100644 docs/source/tubes/server.rst

diff --git a/docs/source/tubes/server.rst b/docs/source/tubes/server.rst
deleted file mode 100644
index d6c6212d8..000000000
--- a/docs/source/tubes/server.rst
+++ /dev/null
@@ -1,10 +0,0 @@
-.. testsetup:: *
-
-   from pwn import *
-   from pwnlib.tubes.server import server
-
-:mod:`pwnlib.tubes.server` --- Server
-==================================================
-
-.. automodule:: pwnlib.tubes.server
-   :members:
diff --git a/docs/source/tubes/sockets.rst b/docs/source/tubes/sockets.rst
index a799ff1bd..ca738b97c 100644
--- a/docs/source/tubes/sockets.rst
+++ b/docs/source/tubes/sockets.rst
@@ -23,3 +23,9 @@
    .. autoclass:: pwnlib.tubes.listen.listen
       :members:
       :show-inheritance:
+
+.. automodule:: pwnlib.tubes.server
+
+   .. autoclass:: pwnlib.tubes.server.server
+      :members:
+      :show-inheritance:

From f189b34e66f4c7bc19d2c66f28e035c32a91a045 Mon Sep 17 00:00:00 2001
From: Devin Neal <devneal17@gmail.com>
Date: Mon, 11 Dec 2017 15:15:59 -0500
Subject: [PATCH 8/8] fix import in sockets.rst

---
 docs/source/tubes/sockets.rst | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/source/tubes/sockets.rst b/docs/source/tubes/sockets.rst
index ca738b97c..bf5967b47 100644
--- a/docs/source/tubes/sockets.rst
+++ b/docs/source/tubes/sockets.rst
@@ -1,6 +1,7 @@
 .. testsetup:: *
 
    from pwn import *
+   from pwnlib.tubes.server import server
 
 :mod:`pwnlib.tubes.sock` --- Sockets
 ===========================================================