Python socket.MSG_WAITALL Examples

def _SendRecv():
    """Communicate with the Developer Shell server socket."""

    port = int(os.getenv(DEVSHELL_ENV, 0))
    if port == 0:
        raise NoDevshellServer()

    sock = socket.socket()
    sock.connect(('localhost', port))

    data = CREDENTIAL_INFO_REQUEST_JSON
    msg = '%s\n%s' % (len(data), data)
    sock.sendall(_to_bytes(msg, encoding='utf-8'))

    header = sock.recv(6).decode()
    if '\n' not in header:
        raise CommunicationError('saw no newline in the first 6 bytes')
    len_str, json_str = header.split('\n', 1)
    to_read = int(len_str) - len(json_str)
    if to_read > 0:
        json_str += sock.recv(to_read, socket.MSG_WAITALL).decode()

    return CredentialInfoResponse(json_str) 

def _SendRecv():
    """Communicate with the Developer Shell server socket."""

    port = int(os.getenv(DEVSHELL_ENV, 0))
    if port == 0:
        raise NoDevshellServer()

    sock = socket.socket()
    sock.connect(('localhost', port))

    data = CREDENTIAL_INFO_REQUEST_JSON
    msg = '%s\n%s' % (len(data), data)
    sock.sendall(_to_bytes(msg, encoding='utf-8'))

    header = sock.recv(6).decode()
    if '\n' not in header:
        raise CommunicationError('saw no newline in the first 6 bytes')
    len_str, json_str = header.split('\n', 1)
    to_read = int(len_str) - len(json_str)
    if to_read > 0:
        json_str += sock.recv(to_read, socket.MSG_WAITALL).decode()

    return CredentialInfoResponse(json_str) 

def _SendRecv():
  """Communicate with the Developer Shell server socket."""

  port = int(os.getenv(DEVSHELL_ENV, 0))
  if port == 0:
    raise NoDevshellServer()

  import socket

  sock = socket.socket()
  sock.connect(('localhost', port))

  data = CREDENTIAL_INFO_REQUEST_JSON
  msg = '%s\n%s' % (len(data), data)
  sock.sendall(msg.encode())

  header = sock.recv(6).decode()
  if '\n' not in header:
    raise CommunicationError('saw no newline in the first 6 bytes')
  len_str, json_str = header.split('\n', 1)
  to_read = int(len_str) - len(json_str)
  if to_read > 0:
    json_str += sock.recv(to_read, socket.MSG_WAITALL).decode()

  return CredentialInfoResponse(json_str) 

def _SendRecv():
    """Communicate with the Developer Shell server socket."""

    port = int(os.getenv(DEVSHELL_ENV, 0))
    if port == 0:
        raise NoDevshellServer()

    sock = socket.socket()
    sock.connect(('localhost', port))

    data = CREDENTIAL_INFO_REQUEST_JSON
    msg = '{0}\n{1}'.format(len(data), data)
    sock.sendall(_helpers._to_bytes(msg, encoding='utf-8'))

    header = sock.recv(6).decode()
    if '\n' not in header:
        raise CommunicationError('saw no newline in the first 6 bytes')
    len_str, json_str = header.split('\n', 1)
    to_read = int(len_str) - len(json_str)
    if to_read > 0:
        json_str += sock.recv(to_read, socket.MSG_WAITALL).decode()

    return CredentialInfoResponse(json_str) 

def _SendRecv():
    """Communicate with the Developer Shell server socket."""

    port = int(os.getenv(DEVSHELL_ENV, 0))
    if port == 0:
        raise NoDevshellServer()

    sock = socket.socket()
    sock.connect(('localhost', port))

    data = CREDENTIAL_INFO_REQUEST_JSON
    msg = '%s\n%s' % (len(data), data)
    sock.sendall(_to_bytes(msg, encoding='utf-8'))

    header = sock.recv(6).decode()
    if '\n' not in header:
        raise CommunicationError('saw no newline in the first 6 bytes')
    len_str, json_str = header.split('\n', 1)
    to_read = int(len_str) - len(json_str)
    if to_read > 0:
        json_str += sock.recv(to_read, socket.MSG_WAITALL).decode()

    return CredentialInfoResponse(json_str) 

def _read_packet(self, socket_):
        if hasattr(socket, "MSG_WAITALL"):
            data = socket_.recv(3, socket.MSG_WAITALL)
        else:
            # Windows lacks MSG_WAITALL
            data = b''
            while len(data) < 3:
                data += socket_.recv(3 - len(data))

        type, length = struct.unpack(">BH", data)
        body_len = length - 3
        while body_len > 0:
            read_len = 32767 if body_len > 32767 else body_len
            data += socket_.recv(read_len, socket.MSG_WAITALL)
            body_len -= read_len
        return type, data 

def _SendRecv():
    """Communicate with the Developer Shell server socket."""

    port = int(os.getenv(DEVSHELL_ENV, 0))
    if port == 0:
        raise NoDevshellServer()

    sock = socket.socket()
    sock.connect(('localhost', port))

    data = CREDENTIAL_INFO_REQUEST_JSON
    msg = '{0}\n{1}'.format(len(data), data)
    sock.sendall(_helpers._to_bytes(msg, encoding='utf-8'))

    header = sock.recv(6).decode()
    if '\n' not in header:
        raise CommunicationError('saw no newline in the first 6 bytes')
    len_str, json_str = header.split('\n', 1)
    to_read = int(len_str) - len(json_str)
    if to_read > 0:
        json_str += sock.recv(to_read, socket.MSG_WAITALL).decode()

    return CredentialInfoResponse(json_str) 

def test_nucleo_usart_debug_write():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect(('127.0.0.1', PORT))


    time.sleep(.1)
    for c in test_string:
        qemu.wm(0x40004404, 1, c)
    reply =  s.recv(len(test_string), socket.MSG_WAITALL)
    assert_equal(reply, test_string)


    time.sleep(.1)
    for c in test_string:
        qemu.wm(0x40004404, 4, c)
    reply =  s.recv(len(test_string), socket.MSG_WAITALL)
    assert_equal(reply, test_string) 

def test_nucleo_usart_read():
    
    #import IPython; IPython.embed()
    
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.connect(('127.0.0.1', PORT))
    qemu.cont()

    data = s.recv(len(test_string), socket.MSG_WAITALL)

    assert_equal(data, test_string) 

def recv(self, timeout=.0001):
        """Receive data of arbitrary length."""
        # First get data length
        self.settimeout(timeout)
        length = self._socket.recv(8)
        length = int.from_bytes(length, "little")
        if length == 0:
            raise ConnectionAbortedError
        # Then receive data
        self.settimeout(None)
        return self._socket.recv(length, socket.MSG_WAITALL) 

def recv_all(self, size):
            if hasattr(socket, "MSG_WAITALL"):
                data = self.request.recv(size, socket.MSG_WAITALL)
            else:
                # Windows lacks MSG_WAITALL
                data = b''
                while len(data) < size:
                    data += self.request.recv(size - len(data))
            return data 

def checkVersion(self, ignoreversioncheck=False, reconnect=True):
        if DEBUG:
            print >> sys.stderr, "checkVersion(reconnect=%s)   ignoreversioncheck=%s" % (reconnect, ignoreversioncheck)
        self.__send('host:version', reconnect=False)
        # HACK: MSG_WAITALL not available on windows
        #version = self.socket.recv(8, socket.MSG_WAITALL)
        version = self.__readExactly(self.socket, 8)

        VALID_ADB_VERSIONS = ["00040020", "0004001f"]

        if not (version in VALID_ADB_VERSIONS) and not ignoreversioncheck:
            raise RuntimeError("ERROR: Incorrect ADB server version %s (expecting one of %s)" % (version, VALID_ADB_VERSIONS))
        if reconnect:
            self.__connect() 

def recv(sock, size):
    """Receives exactly `size` bytes.  This function blocks the thread."""
    data = sock.recv(size, socket.MSG_WAITALL)
    if len(data) < size:
        raise socket.error(ECONNRESET, 'Connection closed')
    return data 

def connect_zabbix(self, payload):
        """
        Method used to send information to Zabbix
        :param payload: refers to the json message prepared to send to Zabbix
        :rtype : returns the response received by the Zabbix API
        """
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.connect((self.zabbix_host, int(self.zabbix_port)))
        s.send(payload)
        # read its response, the first five bytes are the header again
        response_header = s.recv(5, socket.MSG_WAITALL)
        if not response_header == 'ZBXD\1':
            raise ValueError('Got invalid response')

        # read the data header to get the length of the response
        response_data_header = s.recv(8, socket.MSG_WAITALL)
        response_data_header = response_data_header[:4]
        response_len = struct.unpack('i', response_data_header)[0]

        # read the whole rest of the response now that we know the length
        response_raw = s.recv(response_len, socket.MSG_WAITALL)
        s.close()

        response = json.loads(response_raw)

        return response 

def testMsgWaitAllConfig(self):
        if platform.system() == "Windows":
            # default config should be False on these platforms even though socket.MSG_WAITALL might exist
            assert not socketutil.USE_MSG_WAITALL
        else:
            # on all other platforms, default config should be True (as long as socket.MSG_WAITALL exists)
            if hasattr(socket, "MSG_WAITALL"):
                assert socketutil.USE_MSG_WAITALL
            else:
                assert not socketutil.USE_MSG_WAITALL 

def testMsgWaitallProblems2(self):
        class ReceiveThread(threading.Thread):
            def __init__(self, sock, sizes):
                super(ReceiveThread, self).__init__()
                self.sock = sock
                self.sizes = sizes

            def run(self):
                cs, _ = self.sock.accept()
                for size in self.sizes:
                    data = socketutil.receive_data(cs, size)
                    socketutil.send_data(cs, data)
                cs.close()

        ss = socketutil.create_socket(bind=("localhost", 0))
        SIZES = [1000, 10000, 32000, 32768, 32780, 41950, 41952, 42000, 65000, 65535, 65600, 80000, 999999]
        serverthread = ReceiveThread(ss, SIZES)
        serverthread.setDaemon(True)
        serverthread.start()
        port = ss.getsockname()[1]
        cs = socketutil.create_socket(connect=("localhost", port), timeout=2)
        # test some sizes that might be problematic with MSG_WAITALL and check that they work fine
        for size in SIZES:
            socketutil.send_data(cs, b"x" * size)
            data = socketutil.receive_data(cs, size)
            assert size == len(data)
        serverthread.join()
        ss.close()
        cs.close() 

def testMsgWaitallProblems(self):
        ss = socketutil.create_socket(bind=("localhost", 0), timeout=2)
        port = ss.getsockname()[1]
        cs = socketutil.create_socket(connect=("localhost", port), timeout=2)
        a = ss.accept()
        # test some sizes that might be problematic with MSG_WAITALL and check that they work fine
        for size in [1000, 10000, 32000, 32768, 32780, 41950, 41952, 42000, 65000, 65535, 65600, 80000]:
            socketutil.send_data(cs, b"x" * size)
            data = socketutil.receive_data(a[0], size)
            socketutil.send_data(a[0], data)
            data = socketutil.receive_data(cs, size)
            assert size == len(data)
        a[0].close()
        ss.close()
        cs.close() 

def receive_data(sock: socket.socket, size: int) -> bytes:
    """Retrieve a given number of bytes from a socket.
    It is expected the socket is able to supply that number of bytes.
    If it isn't, an exception is raised (you will not get a zero length result
    or a result that is smaller than what you asked for). The partial data that
    has been received however is stored in the 'partialData' attribute of
    the exception object."""
    try:
        delays = __retrydelays()
        msglen = 0
        data = bytearray()
        if USE_MSG_WAITALL and not hasattr(sock, "getpeercert"):    # ssl doesn't support recv flags
            while True:
                try:
                    chunk = sock.recv(size, socket.MSG_WAITALL)
                    if len(chunk) == size:
                        return chunk
                    # less data than asked, drop down into normal receive loop to finish
                    msglen = len(chunk)
                    data.extend(chunk)
                    break
                except socket.timeout:
                    raise TimeoutError("receiving: timeout")
                except socket.error as x:
                    err = getattr(x, "errno", x.args[0])
                    if err not in ERRNO_RETRIES:
                        raise ConnectionClosedError("receiving: connection lost: " + str(x))
                    time.sleep(next(delays))  # a slight delay to wait before retrying
        # old fashioned recv loop, we gather chunks until the message is complete
        while True:
            try:
                while msglen < size:
                    # 60k buffer limit avoids problems on certain OSes like VMS, Windows
                    chunk = sock.recv(min(60000, size - msglen))
                    if not chunk:
                        break
                    data.extend(chunk)
                    msglen += len(chunk)
                if len(data) != size:
                    err = ConnectionClosedError("receiving: not enough data")
                    err.partialData = data  # store the message that was received until now
                    raise err
                return data  # yay, complete
            except socket.timeout:
                raise TimeoutError("receiving: timeout")
            except socket.error as x:
                err = getattr(x, "errno", x.args[0])
                if err not in ERRNO_RETRIES:
                    raise ConnectionClosedError("receiving: connection lost: " + str(x))
                time.sleep(next(delays))  # a slight delay to wait before retrying
    except socket.timeout:
        raise TimeoutError("receiving: timeout") 

def run(self):
        video_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        video_socket.connect((self.host, ardrone.constant.VIDEO_PORT))

        nav_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        nav_socket.setblocking(False)
        nav_socket.bind(('', ardrone.constant.NAVDATA_PORT))
        nav_socket.sendto(b'\x01\x00\x00\x00', (self.host, ardrone.constant.NAVDATA_PORT))

        stopping = False
        while not stopping:
            inputready, outputready, exceptready = select.select([nav_socket, video_socket, self.com_pipe], [], [])
            for i in inputready:
                if i == video_socket:
                    # get first few bytes of header
                    data = video_socket.recv(12, socket.MSG_WAITALL)
                    if len(data) != 12:
                        continue

                    # decode relevant portions of the header
                    sig_p, sig_a, sig_v, sig_e, version, codec, header, payload = struct.unpack('4cBBHI', data)

                    # check signature (and ignore packet otherwise)
                    if sig_p != b'P' or sig_a != b'a' or sig_v != b'V' or sig_e != b'E':
                        continue

                    # get remaining frame
                    data += video_socket.recv(header - 12 + payload, socket.MSG_WAITALL)

                    try:
                        # decode the frame
                        image = ardrone.video.decode(data)
                        self.video_pipe.send(image)
                    except ardrone.video.DecodeError:
                        pass
                elif i == nav_socket:
                    while 1:
                        try:
                            data = nav_socket.recv(65535)
                        except IOError:
                            # we consumed every packet from the socket and
                            # continue with the last one
                            break
                    navdata = ardrone.navdata.decode(data)
                    self.nav_pipe.send(navdata)
                elif i == self.com_pipe:
                    _ = self.com_pipe.recv()
                    stopping = True
                    break
        video_socket.close()
        nav_socket.close() 

def from_file(input_file, logger=None):
        """This static method acts as a constructor and returns an input
        packet with the proper class, based on the packet headers.
        The "input_file" parameter must either be a file or a sockect object.

        """

        if isinstance(input_file, _socketobject):
            def read_file(count):
                return input_file.recv(count, MSG_WAITALL)
        elif hasattr(file, "read") and callable(file.read):
            def read_file(count):
                return input_file.read(count)
        else:
            raise ValueError("'input_file' must either be a socket object or"
                             " provide a 'read' method")

        fields = ("rpc_vers", "rpc_vers_minor", "ptype", "pfc_flags", "drep",
                  "frag_length", "auth_length", "call_id")

        header_data = read_file(16)
        # len_header_data = len(header_data)
        # if len_header_data < 16:
        #     raise RTSParsingException("read only %d header bytes from input"
        #                               % len_header_data)

        # TODO: value validation
        values = unpack_from("<bbbblhhl", header_data)
        if values[2] == DCERPC_PKT_FAULT:
            packet_class = RPCFaultPacket
        elif values[2] == DCERPC_PKT_BIND_ACK:
            packet_class = RPCBindACKPacket
        elif values[2] == DCERPC_PKT_BIND_NAK:
            packet_class = RPCBindNAKPacket
        elif values[2] == DCERPC_PKT_RTS:
            packet_class = RPCRTSPacket
        else:
            packet_class = RPCPacket
        body_data = read_file(values[5] - 16)
        # len_body_data = len(body_data)
        # if len_body_data < (values[5] - 16):
        #     raise RTSParsingException("read only %d body bytes from input"
        #                               % len_body_data)

        packet = packet_class(header_data + body_data, logger)
        packet.header = dict(zip(fields, values))
        packet.offset = 16
        packet.size = values[5]
        packet.parse()

        return packet 

def connect(self):
        if self.sock:
            self.close()

        s = socket.socket()
        s.connect(self.endpoint)

        hello = s.recv(8, socket.MSG_WAITALL)
        if hello != b"FLUX0003":
            raise NotSupportError(self.model_id, self.version)

        self.sock = ssl.SSLSocket(s)

        # Stage 1: Recv randbytes
        self.randbytes = self.recv_bytes(64)

        # Stage 2: Send public key
        strkey = self.client_key.public_key_pem.decode("ascii")
        self.send_text(strkey)

        # Stage 3: Get public key status
        resp = self.recv_text()

        if resp == "sign":
            # Stage 4.a: Sign
            doc = HMAC(self.uuid.bytes, self.randbytes, sha1).digest()
            signature = self.client_key.sign(doc)
            self.send_text(to_hex(signature))

        elif resp == "password":
            # Stage 4.b: Send password
            return

        elif resp.startswith("error "):
            raise raise_error(resp)
        else:
            # raise NotSupportError("Auth method %s not support", resp)
            raise NotSupportError(self.model_id, self.version)

        resp = self.recv_text()
        if resp == "ok":
            self._authorized = True
            return
        elif resp.startswith("error "):
            err = resp[6:]
            if err == "AUTH_ERROR":
                raise AuthError()
            else:
                raise ManagerError(err)