Python socket.MSG_PEEK Examples

def readline_p(self, size=-1):
        """Reads a line from this file"""
        self.__checkMode('r')
        if size < 0:
            size = 20000

        # The integration of the debugger client event loop and the connection
        # to the debugger relies on the two lines of the debugger command being
        # delivered as two separate events. Therefore we make sure we only
        # read a line at a time.
        line = self.sock.recv(size, socket.MSG_PEEK)
        eol = line.find(b'\n')

        if eol >= 0:
            size = eol + 1
        else:
            size = len(line)

        # Now we know how big the line is, read it for real.
        return self.sock.recv(size).decode('utf8', 'backslashreplace') 

def _get_read_size(self, sock, recv_buffer_size, up):
        if self._overhead == 0:
            return recv_buffer_size
        buffer_size = len(sock.recv(recv_buffer_size, socket.MSG_PEEK))
        frame_size = self._tcp_mss - self._overhead
        if up:
            buffer_size = min(buffer_size, self._recv_u_max_size)
            self._recv_u_max_size = min(self._recv_u_max_size + frame_size, BUF_SIZE)
        else:
            buffer_size = min(buffer_size, self._recv_d_max_size)
            self._recv_d_max_size = min(self._recv_d_max_size + frame_size, BUF_SIZE)
        if buffer_size == recv_buffer_size:
            return buffer_size
        if buffer_size > frame_size:
            buffer_size = int(buffer_size / frame_size) * frame_size
        return buffer_size 

def recv(self, bufsiz, flags=None):
        """
        Receive data on the connection.

        :param bufsiz: The maximum number of bytes to read
        :param flags: (optional) The only supported flag is ``MSG_PEEK``,
            all other flags are ignored.
        :return: The string read from the Connection
        """
        buf = _no_zero_allocator("char[]", bufsiz)
        if flags is not None and flags & socket.MSG_PEEK:
            result = _lib.SSL_peek(self._ssl, buf, bufsiz)
        else:
            result = _lib.SSL_read(self._ssl, buf, bufsiz)
        self._raise_ssl_error(self._ssl, result)
        return _ffi.buffer(buf, result)[:] 

def _get_read_size(self, sock, recv_buffer_size, up):
        if self._overhead == 0:
            return recv_buffer_size
        buffer_size = len(sock.recv(recv_buffer_size, socket.MSG_PEEK))
        if up:
            buffer_size = min(buffer_size, self._recv_u_max_size)
            self._recv_u_max_size = min(self._recv_u_max_size + self._tcp_mss - self._overhead, BUF_SIZE)
        else:
            buffer_size = min(buffer_size, self._recv_d_max_size)
            self._recv_d_max_size = min(self._recv_d_max_size + self._tcp_mss - self._overhead, BUF_SIZE)
        if buffer_size == recv_buffer_size:
            return buffer_size
        frame_size = self._tcp_mss - self._overhead
        if buffer_size > frame_size:
            buffer_size = int(buffer_size / frame_size) * frame_size
        return buffer_size 

def _get_read_size(self, sock, recv_buffer_size, up):
        if self._overhead == 0:
            return recv_buffer_size
        buffer_size = len(sock.recv(recv_buffer_size, socket.MSG_PEEK))
        frame_size = self._tcp_mss - self._overhead
        if up:
            buffer_size = min(buffer_size, self._recv_u_max_size)
            self._recv_u_max_size = min(self._recv_u_max_size + frame_size, BUF_SIZE)
        else:
            buffer_size = min(buffer_size, self._recv_d_max_size)
            self._recv_d_max_size = min(self._recv_d_max_size + frame_size, BUF_SIZE)
        if buffer_size == recv_buffer_size:
            return buffer_size
        if buffer_size > frame_size:
            buffer_size = int(buffer_size / frame_size) * frame_size
        return buffer_size 

def _get_read_size(self, sock, recv_buffer_size, up):
        if self._overhead == 0:
            return recv_buffer_size
        buffer_size = len(sock.recv(recv_buffer_size, socket.MSG_PEEK))
        if up:
            buffer_size = min(buffer_size, self._recv_u_max_size)
            self._recv_u_max_size = min(self._recv_u_max_size + self._tcp_mss - self._overhead, BUF_SIZE)
        else:
            buffer_size = min(buffer_size, self._recv_d_max_size)
            self._recv_d_max_size = min(self._recv_d_max_size + self._tcp_mss - self._overhead, BUF_SIZE)
        if buffer_size == recv_buffer_size:
            return buffer_size
        frame_size = self._tcp_mss - self._overhead
        if buffer_size > frame_size:
            buffer_size = int(buffer_size / frame_size) * frame_size
        return buffer_size 

def recv(self, bufsiz, flags=None):
        """
        Receive data on the connection.

        :param bufsiz: The maximum number of bytes to read
        :param flags: (optional) The only supported flag is ``MSG_PEEK``,
            all other flags are ignored.
        :return: The string read from the Connection
        """
        buf = _no_zero_allocator("char[]", bufsiz)
        if flags is not None and flags & socket.MSG_PEEK:
            result = _lib.SSL_peek(self._ssl, buf, bufsiz)
        else:
            result = _lib.SSL_read(self._ssl, buf, bufsiz)
        self._raise_ssl_error(self._ssl, result)
        return _ffi.buffer(buf, result)[:] 

def _get_read_size(self, sock, recv_buffer_size, up):
        if self._overhead == 0:
            return recv_buffer_size
        buffer_size = len(sock.recv(recv_buffer_size, socket.MSG_PEEK))
        frame_size = self._tcp_mss - self._overhead
        if up:
            buffer_size = min(buffer_size, self._recv_u_max_size)
            self._recv_u_max_size = min(self._recv_u_max_size + frame_size, BUF_SIZE)
        else:
            buffer_size = min(buffer_size, self._recv_d_max_size)
            self._recv_d_max_size = min(self._recv_d_max_size + frame_size, BUF_SIZE)
        if buffer_size == recv_buffer_size:
            return buffer_size
        if buffer_size > frame_size:
            buffer_size = int(buffer_size / frame_size) * frame_size
        return buffer_size 

def get_ssl_version(sock):
        # Seth behaves differently depeding on the TLS protocol
        # https://bugs.python.org/issue31453
        # This is an ugly hack (as if the rest of this wasn't...)
    versions = [
        ssl.PROTOCOL_TLSv1,
        ssl.PROTOCOL_TLSv1_1,
        ssl.PROTOCOL_TLSv1_2,
        ]
    firstbytes = sock.recv(16, socket.MSG_PEEK)
    try:
        return versions[firstbytes[10]-1]
    except IndexError:
        print("Unexpected SSL version: %s" % hexlify(firstbytes))
        return versions[-1]


#  def launch_rdp_client():
#      time.sleep(1)
#      p = subprocess.Popen(
#          ["xfreerdp",
#           "/v:%s:%d" % (args.bind_ip, consts.RELAY_PORT),
#           "/u:%s\\%s" % (domain, user),
#          ],
#      ) 

def handle_one_request(self):
        if not self.disable_transport_ssl and self.scheme == 'http':
            leadbyte = self.connection.recv(1, socket.MSG_PEEK)
            if leadbyte in ('\x80', '\x16'):
                server_name = ''
                if leadbyte == '\x16':
                    for _ in xrange(2):
                        leaddata = self.connection.recv(1024, socket.MSG_PEEK)
                        if is_clienthello(leaddata):
                            try:
                                server_name = extract_sni_name(leaddata)
                            finally:
                                break
                try:
                    certfile = CertUtil.get_cert(server_name or 'www.google.com')
                    ssl_sock = ssl.wrap_socket(self.connection, ssl_version=self.ssl_version, keyfile=certfile, certfile=certfile, server_side=True)
                except StandardError as e:
                    if e.args[0] not in (errno.ECONNABORTED, errno.ECONNRESET):
                        logging.exception('ssl.wrap_socket(self.connection=%r) failed: %s', self.connection, e)
                    return
                self.connection = ssl_sock
                self.rfile = self.connection.makefile('rb', self.bufsize)
                self.wfile = self.connection.makefile('wb', 0)
                self.scheme = 'https'
        return BaseHTTPServer.BaseHTTPRequestHandler.handle_one_request(self) 

def handle_tcp_default(sk, dstport):
	# Attempt to guess protocol according to what the client sends
	data = ''
	try:
		rlist, _, _ = select.select([sk], [], [], 30)
		if len(rlist) != 0:
			data = sk.recv(20, socket.MSG_PEEK)
	except Exception as err:
		#print(traceback.format_exc())
		pass

	if data[:3] in SSL_CLIENT_HELLO_SIGNATURES:
		print colored("Guessing this is a SSL/TLS connection, attempting to handshake.", 'red', attrs=['bold'])
		handle_tcp_hexdump_ssl(sk, dstport)
	elif data.startswith("GET "):
		handle_tcp_http(sk, dstport)
	elif data.startswith("CONNECT "):
		handle_tcp_httpproxy(sk, dstport)
	else:
		handle_tcp_hexdump(sk, dstport)
	sk.close()

# UDP DISPATCHER 

def _get_read_size(self, sock, recv_buffer_size, up):
        if self._overhead == 0:
            return recv_buffer_size
        buffer_size = len(sock.recv(recv_buffer_size, socket.MSG_PEEK))
        frame_size = self._tcp_mss - self._overhead
        if up:
            buffer_size = min(buffer_size, self._recv_u_max_size)
            self._recv_u_max_size = min(self._recv_u_max_size + frame_size, BUF_SIZE)
        else:
            buffer_size = min(buffer_size, self._recv_d_max_size)
            self._recv_d_max_size = min(self._recv_d_max_size + frame_size, BUF_SIZE)
        if buffer_size == recv_buffer_size:
            return buffer_size
        if buffer_size > frame_size:
            buffer_size = int(buffer_size / frame_size) * frame_size
        return buffer_size 

def recv(self, bufsiz, flags=None):
        """
        Receive data on the connection.

        :param bufsiz: The maximum number of bytes to read
        :param flags: (optional) The only supported flag is ``MSG_PEEK``,
            all other flags are ignored.
        :return: The string read from the Connection
        """
        buf = _ffi.new("char[]", bufsiz)
        if flags is not None and flags & socket.MSG_PEEK:
            result = _lib.SSL_peek(self._ssl, buf, bufsiz)
        else:
            result = _lib.SSL_read(self._ssl, buf, bufsiz)
        self._raise_ssl_error(self._ssl, result)
        return _ffi.buffer(buf, result)[:] 

def _Dynamic_Receive(self, request, response):
    state = self._LookupSocket(request.socket_descriptor())
    with state.mutex:
      state.SetTimeout(request.timeout_seconds())
      flags = 0
      if request.flags() & remote_socket_service_pb.ReceiveRequest.MSG_PEEK:
        flags |= socket.MSG_PEEK
      received_from = None
      if state.protocol == socket.SOCK_DGRAM:
        data, received_from = state.sock.recvfrom(request.data_size(), flags)
      else:
        data = state.sock.recv(request.data_size(), flags)
      response.set_data(data)
      if received_from:
        self._AddressPortTupleToProto(state.family, received_from,
                                      response.mutable_received_from()) 

def testRecvmsgPeek(self):
        # Check that MSG_PEEK in flags enables examination of pending
        # data without consuming it.

        # Receive part of data with MSG_PEEK.
        msg, ancdata, flags, addr = self.doRecvmsg(self.serv_sock,
                                                   len(MSG) - 3, 0,
                                                   socket.MSG_PEEK)
        self.assertEqual(msg, MSG[:-3])
        self.checkRecvmsgAddress(addr, self.cli_addr)
        self.assertEqual(ancdata, [])
        # Ignoring MSG_TRUNC here (so this test is the same for stream
        # and datagram sockets).  Some wording in POSIX seems to
        # suggest that it needn't be set when peeking, but that may
        # just be a slip.
        self.checkFlags(flags, eor=False,
                        ignore=getattr(socket, "MSG_TRUNC", 0))

        # Receive all data with MSG_PEEK.
        msg, ancdata, flags, addr = self.doRecvmsg(self.serv_sock,
                                                   len(MSG), 0,
                                                   socket.MSG_PEEK)
        self.assertEqual(msg, MSG)
        self.checkRecvmsgAddress(addr, self.cli_addr)
        self.assertEqual(ancdata, [])
        self.checkFlags(flags, eor=True)

        # Check that the same data can still be received normally.
        msg, ancdata, flags, addr = self.doRecvmsg(self.serv_sock, len(MSG))
        self.assertEqual(msg, MSG)
        self.checkRecvmsgAddress(addr, self.cli_addr)
        self.assertEqual(ancdata, [])
        self.checkFlags(flags, eor=True) 

def testPeek(self):
        data, addr = self.serv.recvfrom(self.bufsize, socket.MSG_PEEK)
        self.assertEqual(self.data, data)
        data, addr = self.serv.recvfrom(self.bufsize)
        self.assertEqual(self.data, data) 

def recv(self, size, flag=0):
        if flag & socket.MSG_PEEK > 0:
            raise RobotError("recv flag not support",
                             error_symbol=("BAD_PARAMS",
                                           "MSG_PEEK_NOT_ALLOWED"))

        if self._decoder:
            buf = self._sock.recv(size, flag)
            return self._decoder.decrypt(buf)
        else:
            raise RobotNotReadyError("Handshake not complete") 

def test_peek(self):
        """
        `Connection.recv` peeks into the connection if `socket.MSG_PEEK` is
        passed.
        """
        server, client = loopback()
        server.send(b'xy')
        assert client.recv(2, MSG_PEEK) == b'xy'
        assert client.recv(2, MSG_PEEK) == b'xy'
        assert client.recv(2) == b'xy' 

def recv_into(self, buffer, nbytes=None, flags=None):
        """
        Receive data on the connection and copy it directly into the provided
        buffer, rather than creating a new string.

        :param buffer: The buffer to copy into.
        :param nbytes: (optional) The maximum number of bytes to read into the
            buffer. If not present, defaults to the size of the buffer. If
            larger than the size of the buffer, is reduced to the size of the
            buffer.
        :param flags: (optional) The only supported flag is ``MSG_PEEK``,
            all other flags are ignored.
        :return: The number of bytes read into the buffer.
        """
        if nbytes is None:
            nbytes = len(buffer)
        else:
            nbytes = min(nbytes, len(buffer))

        # We need to create a temporary buffer. This is annoying, it would be
        # better if we could pass memoryviews straight into the SSL_read call,
        # but right now we can't. Revisit this if CFFI gets that ability.
        buf = _no_zero_allocator("char[]", nbytes)
        if flags is not None and flags & socket.MSG_PEEK:
            result = _lib.SSL_peek(self._ssl, buf, nbytes)
        else:
            result = _lib.SSL_read(self._ssl, buf, nbytes)
        self._raise_ssl_error(self._ssl, result)

        # This strange line is all to avoid a memory copy. The buffer protocol
        # should allow us to assign a CFFI buffer to the LHS of this line, but
        # on CPython 3.3+ that segfaults. As a workaround, we can temporarily
        # wrap it in a memoryview.
        buffer[:result] = memoryview(_ffi.buffer(buf, result))

        return result 

def recv(self, x=MTU):
        pkt = self.ins.recv(x, socket.MSG_PEEK)
        x = len(pkt)
        if x == 0:
            raise socket.error((100,"Underlying stream socket tore down"))
        pkt = self.basecls(pkt)
        pad = pkt.getlayer(conf.padding_layer)
        if pad is not None and pad.underlayer is not None:
            del(pad.underlayer.payload)
        while pad is not None and not isinstance(pad, NoPayload):
            x -= len(pad.load)
            pad = pad.payload
        self.ins.recv(x)
        return pkt 

def test_peek(self):
        server, client = loopback()
        server.send(b'xy')

        for _ in range(2):
            output_buffer = bytearray(5)
            assert client.recv_into(output_buffer, flags=MSG_PEEK) == 2
            assert output_buffer == bytearray(b'xy\x00\x00\x00') 

def recv(self, x=MTU):
        pkt = self.ins.recv(x, socket.MSG_PEEK)
        x = len(pkt)
        if x == 0:
            raise socket.error((100,"Underlying stream socket tore down"))
        pkt = self.basecls(pkt)
        pad = pkt.getlayer(conf.padding_layer)
        if pad is not None and pad.underlayer is not None:
            del(pad.underlayer.payload)
        while pad is not None and not isinstance(pad, NoPayload):
            x -= len(pad.load)
            pad = pad.payload
        self.ins.recv(x)
        return pkt 

def recv(self, x=MTU):
        pkt = self.ins.recv(x, socket.MSG_PEEK)
        x = len(pkt)
        if x == 0:
            raise socket.error((100,"Underlying stream socket tore down"))
        pkt = self.basecls(pkt)
        pad = pkt.getlayer(conf.padding_layer)
        if pad is not None and pad.underlayer is not None:
            del(pad.underlayer.payload)
        while pad is not None and not isinstance(pad, NoPayload):
            x -= len(pad.load)
            pad = pad.payload
        self.ins.recv(x)
        return pkt 

def recv(self, x=MTU):
        pkt = self.ins.recv(x, socket.MSG_PEEK)
        x = len(pkt)
        if x == 0:
            raise socket.error((100,"Underlying stream socket tore down"))
        pkt = self.basecls(pkt)
        pad = pkt.getlayer(conf.padding_layer)
        if pad is not None and pad.underlayer is not None:
            del(pad.underlayer.payload)
        while pad is not None and not isinstance(pad, NoPayload):
            x -= len(pad.load)
            pad = pad.payload
        self.ins.recv(x)
        return pkt 

def handle(self):
        self.__class__.handle_num += 1

        if self.try_redirect():
            return

        sockets = [self.conn]
        try:
            r, w, e = select.select(sockets, [], [])
            socks_version = self.conn.recv(1, socket.MSG_PEEK)
            if not socks_version:
                return

            if socks_version == b"\x04":
                self.socks4_handler()
            elif socks_version == b"\x05":
                self.socks5_handler()
            elif socks_version == b"C":
                self.https_handler()
            elif socks_version in [b"G", b"P", b"D", b"O", b"H", b"T"]:
                self.http_handler()
            else:
                xlog.warn("socks version:%s[%s] not supported", socks_version, utils.str2hex(socks_version))
                return

        except socket.error as e:
            xlog.warn('socks handler read error:%r', e)
        except Exception as e:
            xlog.exception("any err:%r", e) 

def http_handler(self):
        req_data = self.conn.recv(65537, socket.MSG_PEEK)
        rp = req_data.split(b"\r\n")
        req_line = rp[0]

        words = req_line.split()
        if len(words) == 3:
            method, url, http_version = words
        elif len(words) == 2:
            method, url = words
            http_version = b"HTTP/1.1"
        else:
            xlog.warn("http req line fail:%s", req_line)
            return

        if url.lower().startswith(b"http://"):
            o = urllib.parse.urlparse(url)
            host, port = netloc_to_host_port(o.netloc)

            url_prex_len = url[7:].find(b"/")
            if url_prex_len >= 0:
                url_prex_len += 7
                path = url[url_prex_len:]
            else:
                url_prex_len = len(url)
                path = b"/"
        else:
            # not proxy request, should be PAC
            xlog.debug("PAC %s %s from:%s", method, url, self.client_address)
            handler = pac_server.PacHandler(self.conn, self.client_address, None, xlog)
            return handler.handle()

        sock = SocketWrap(self.conn, self.client_address[0], self.client_address[1])
        sock.replace_pattern = [url[:url_prex_len], ""]

        xlog.debug("http %r connect to %s:%d %s %s", self.client_address, host, port, method, path)
        handle_domain_proxy(sock, host, port, self.client_address) 

def do_CONNECT(self):
        """deploy fake cert to client"""
        host, _, port = self.path.rpartition(b':')
        port = int(port)
        if port not in (80, 443):
            xlog.warn("CONNECT %s port:%d not support", host, port)
            return

        certfile = CertUtil.get_cert(host)
        self.wfile.write(b'HTTP/1.1 200 Connection Established\r\n\r\n')
        self.wfile.flush()
        #self.conntunnel = True
 
        leadbyte = self.connection.recv(1, socket.MSG_PEEK)
        if leadbyte in (b'\x80', b'\x16'):
            try:
                ssl_sock = ssl.wrap_socket(self.connection, keyfile=CertUtil.cert_keyfile, certfile=certfile, server_side=True)
            except ssl.SSLError as e:
                xlog.info('ssl error: %s, create full domain cert for host:%s', e, host)
                certfile = CertUtil.get_cert(host, full_name=True)
                return
            except Exception as e:
                if e.args[0] not in (errno.ECONNABORTED, errno.ECONNRESET):
                    xlog.exception('ssl.wrap_socket(self.connection=%r) failed: %s path:%s, errno:%s', self.connection, e, self.path, e.args[0])
                return

            self.__realwfile = self.wfile
            self.__realrfile = self.rfile
            self.connection = ssl_sock
            self.rfile = self.connection.makefile('rb', self.bufsize)
            self.wfile = self.connection.makefile('wb', 0)

        self.close_connection = 0 

def serve():
    bindsocket = socket.socket()
    bindsocket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    #bindsocket.bind(('localhost', PORT))
    bindsocket.bind(('', PORT))
    bindsocket.listen(5)

    print("serving on port", PORT)

    while True:
        try:
            newsocket, from_addr = bindsocket.accept()
            peek = newsocket.recv(1024, socket.MSG_PEEK)
            if peek.startswith("\x16"):
                connstream = ssl.wrap_socket(
                        newsocket,
                        server_side=True,
                        certfile='self.pem',
                        ssl_version=ssl.PROTOCOL_TLSv1)
            else:
                connstream = newsocket

            do_request(connstream, from_addr)

        except Exception:
            traceback.print_exc() 

def recv_into(self, buffer, nbytes=None, flags=None):
        """
        Receive data on the connection and copy it directly into the provided
        buffer, rather than creating a new string.

        :param buffer: The buffer to copy into.
        :param nbytes: (optional) The maximum number of bytes to read into the
            buffer. If not present, defaults to the size of the buffer. If
            larger than the size of the buffer, is reduced to the size of the
            buffer.
        :param flags: (optional) The only supported flag is ``MSG_PEEK``,
            all other flags are ignored.
        :return: The number of bytes read into the buffer.
        """
        if nbytes is None:
            nbytes = len(buffer)
        else:
            nbytes = min(nbytes, len(buffer))

        # We need to create a temporary buffer. This is annoying, it would be
        # better if we could pass memoryviews straight into the SSL_read call,
        # but right now we can't. Revisit this if CFFI gets that ability.
        buf = _no_zero_allocator("char[]", nbytes)
        if flags is not None and flags & socket.MSG_PEEK:
            result = _lib.SSL_peek(self._ssl, buf, nbytes)
        else:
            result = _lib.SSL_read(self._ssl, buf, nbytes)
        self._raise_ssl_error(self._ssl, result)

        # This strange line is all to avoid a memory copy. The buffer protocol
        # should allow us to assign a CFFI buffer to the LHS of this line, but
        # on CPython 3.3+ that segfaults. As a workaround, we can temporarily
        # wrap it in a memoryview.
        buffer[:result] = memoryview(_ffi.buffer(buf, result))

        return result 

def recv(self, x=MTU):
        pkt = self.ins.recv(x, socket.MSG_PEEK)
        x = len(pkt)
        if x == 0:
            raise socket.error((100,"Underlying stream socket tore down"))
        pkt = self.basecls(pkt)
        pad = pkt.getlayer(Padding)
        if pad is not None and pad.underlayer is not None:
            del(pad.underlayer.payload)
        while pad is not None and not isinstance(pad, NoPayload):
            x -= len(pad.load)
            pad = pad.payload
        self.ins.recv(x)
        return pkt