def connect(self, address):
"""
Connect the socket to a remote address. For IP sockets, the address
is a pair (host, port).
Please note: only support SOCK_RE mode,
A socket may only connect if it's unbound (clean socket)
:param address: YO remote address ('A.B')
"""
with self.main_lock:
if self.__protocol != SOCK_RE:
raise YOREUnsupportedProtocolAction("Unexpected cmd for this socket. Use with SOCK_RE",
SOCK_RE, self.__protocol)
if self.__state != STATE.NONE:
raise YOREUnsupportedProtocolAction("Error, must be in STATE %d to connect" % STATE.NONE,
STATE.NONE, self.__state)
host, port = address
self.__yo_dest = host
if type(host) is str:
a, b = host.split(".")
self.__yo_dest = struct.unpack(">H", struct.pack(">BB", int(a), int(b)))[0]
self.re_port = port
args = [self.__yo_dest, self.re_port]
return self.__safe_send_and_recv_ipc(SERIALIZER_CMD.CONNECT, STATE.NONE, args,
upgrade_state_if_successful=STATE.CONNECTION_ESTABLISHED)[0]
def recvfrom(self, buffer_size):
"""
Receive From command. This function is blocking until there's data to receive.
Notice that like real sockets, stating buffer_size does not necessarily mean the whole
buffer is going to get filled. When the ProtocolDaemon has data to send back, it will do so.
The buffer_size indicates the maximum data size that the user is able to chew at the moment.
NON-BLOCKING MODE: May return false if disconnected
BLOCKING MODE: May raise YOREException if disconnected
Please note: only support SOCK_RAW mode
:param buffer_size: max buffer length to receive
:return: (address, buffer)
"""
with self.main_lock:
if self.__protocol != SOCK_RAW:
raise YOREUnsupportedProtocolAction("Unexpected cmd for this socket. Use with SOCK_RAW",
SOCK_RAW, self.__protocol)
(res, data) = self.__safe_send_and_recv_ipc(SERIALIZER_CMD.RECVFROM, STATE.BOUND,
[buffer_size], return_command_types=(SERIALIZER_CMD.RECVFROM,
SERIALIZER_CMD.RECV_EMPTY,
SERIALIZER_CMD.CLOSE))
if data[0] == SERIALIZER_CMD.RECV_EMPTY:
if self.__blocking == BLOCKING:
raise YORETimeoutException
else:
return False
res = ((data[1], data[2]) if data[0] == SERIALIZER_CMD.RECVFROM else False) if res else False
# addition - add exception
if self.__blocking == BLOCKING and res == False:
raise YOREException
return res
def __safe_send_and_recv_ipc(self, command_type, should_be_in_state, args, upgrade_state_if_successful=None,
return_command_types=None):
"""
send and recv ipc with many security checks according to the parameters.
for example: to send ipc of "SEND" command we will send SERIALIZER_CMD.SEND and expect to get
reply with the same enum. Also we demand that our state will be STATE.CONNECTION_ESTABLISHED
can raise YOREUnsupportedProtocolAction, YOREDaemonToAPIValueError
:param command_type: SERIALIZER_CMD enum - the cmd to send
:param should_be_in_state: STATE enum - what self.state should be
:param args: [data] - to send to the ipc
:param upgrade_state_if_successful: STATE enum - if everything goes fine - change self.state to new one
:param return_command_types: SERIALIZER_CMD enum list - what to expect to receive -
- if None will use the command_type instead
:return: (True, [data]) or (False,None)
"""
if return_command_types is None:
return_command_types = [command_type]
# Failure check
if should_be_in_state is not None and self.__state != should_be_in_state:
raise YOREUnsupportedProtocolAction("Error, must be in STATE %d to connect" % should_be_in_state,
should_be_in_state, self.__state)
# The real job is in here - send the ipc to the daemon and receive an answer
result, data = self.__send_and_recv_ipc(command_type, args)
# Failure checks
if result == self.__serializer.RESULT_ERROR:
if self.__state != STATE.NONE:
self.__state = STATE.NONE
self.__ipc.close()
if len(data) == 3:
raise YOREProtocolError(data[1], data[2])
elif len(data) == 2:
raise YOREProtocolError(data[1])
else:
raise YOREProtocolError()
elif result != self.__serializer.RESULT_SUCCESS:
self.__state = STATE.NONE
self.__ipc.close()
raise YOREDaemonToAPIValueError(
"Got an unexpected cmd result %d for cmd %d. closing socket!" % (result, command_type),
result, command_type)
if command_type not in return_command_types:
self.__state = STATE.NONE
self.__ipc.close()
raise YOREDaemonToAPIValueError(
"Got an unexpected cmd type %d for cmd %d. closing socket!" % (data[0], command_type),
[0], command_type)
# If we got to here - success!
if upgrade_state_if_successful is not None:
self.__state = upgrade_state_if_successful
return True, data
def listen(self, backlog=1):
"""
Socket listen command - listens for new connections.
Please note: only support SOCK_RE mode,
A socket may only listen if it's bound to an address
:param backlog: At the moment, we don't care about listen "backlog" functionality, and
we only maintain this parameter to keep the listen identical to a normal
socket's listen.
"""
with self.main_lock:
if self.__protocol != SOCK_RE:
raise YOREUnsupportedProtocolAction("Unexpected cmd for this socket. Use with SOCK_RE",
SOCK_RE, self.__protocol)
return self.__safe_send_and_recv_ipc(SERIALIZER_CMD.LISTEN, STATE.BOUND, [backlog],
upgrade_state_if_successful=STATE.LISTEN)[0]
def accept(self):
"""
Accept command.
On success, This command returns a new Resocket (initialized with the
correct YO address, destination port and session id), and the YO
Please note: only support SOCK_RE mode
"""
with self.main_lock:
if self.__protocol != SOCK_RE:
raise YOREUnsupportedProtocolAction("Unexpected cmd for this socket. Use with SOCK_RE",
SOCK_RE, self.__protocol)
(res, data) = self.__safe_send_and_recv_ipc(SERIALIZER_CMD.ACCEPT, STATE.LISTEN, [])
if not res:
return None
yo_addr, self.__yo_dest, re_port, self.__sid, last_packet = data[1:6]
accept_parameters = [self.__yo_dest, re_port, self.__sid, last_packet, self.__blocking, self.__timeout]
return Resocket(SOCK_RE, accept_params=accept_parameters, daemon_port=self.daemon_port), self.__yo_dest
def send(self, send_data):
"""
Send data command.
May raise YOREException if disconnected
Please note: only support SOCK_RE mode
:param send_data: buffer as string. if is empty - ignores and return None
"""
with self.main_lock:
if send_data is None or send_data == '':
return None
if self.__protocol != SOCK_RE:
raise YOREUnsupportedProtocolAction("Unexpected cmd for this socket. Use with SOCK_RE",
SOCK_RE, self.__protocol)
res = self.__safe_send_and_recv_ipc(SERIALIZER_CMD.SEND, STATE.CONNECTION_ESTABLISHED, [send_data])[0]
# addition - add exception
if res == False or res is None:
raise YOREException
return res
def sendto(self, address, send_data):
"""
Sendto data command
May raise YOREException if disconnected
Please note: only support SOCK_RAW mode
:param address: YO remote address ('A.B')
:param send_data: buffer as string. if is empty - ignores and return None
"""
with self.main_lock:
if send_data is None or send_data == '':
return None
if self.__protocol != SOCK_RAW:
raise YOREUnsupportedProtocolAction("Unexpected cmd for this socket. Use with SOCK_RAW",
SOCK_RAW, self.__protocol)
res = self.__safe_send_and_recv_ipc(SERIALIZER_CMD.SENDTO, None, [address, send_data])[0]
# addition - add exception
if res == False or res is None:
raise YOREException
return res