def _read_stream(self, stream: asyncio.StreamReader) -> None:
"""
:param stream: asyncio.StreamReader
:return: None
"""
while not stream.at_eof():
data = yield from stream.read(1)
if data == b'\n':
self.__new_element()
else:
self.__update_last_item_text(data)
def __init__(self, loop, config):
super().__init__(loop=loop)
self._loop = loop
self._config = config
self._transport = None
self._write_pending_data_topic = [] # tuple (data, topic)
self._connected = False
self._encryptor = cryptor.Cryptor(self._config['password'], self._config['method'])
self._peername = None
self._reader_task = None
self._data_task = None
self._keepalive_task = None
self._keepalive_timeout = self._config['timeout']
self._reader_ready = None
self._reader_stopped = asyncio.Event(loop=self._loop)
self._stream_reader = StreamReader(loop=self._loop)
self._stream_writer = None
self._reader = None
self._topic_to_clients = {}
self._queue = Queue(loop=loop)
async def _handshake(
initiator: "HandshakeInitiator",
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
token: CancelToken,
) -> Tuple[bytes, bytes, keccak_256, keccak_256]:
"""See the handshake() function above.
This code was factored out into this helper so that we can create Peers with directly
connected readers/writers for our tests.
"""
initiator_nonce = keccak(os.urandom(HASH_LEN))
auth_msg = initiator.create_auth_message(initiator_nonce)
auth_init = initiator.encrypt_auth_message(auth_msg)
writer.write(auth_init)
auth_ack = await token.cancellable_wait(
reader.read(ENCRYPTED_AUTH_ACK_LEN), timeout=REPLY_TIMEOUT)
if reader.at_eof():
# This is what happens when Parity nodes have blacklisted us
# (https://github.com/ethereum/py-evm/issues/901).
raise HandshakeFailure("%s disconnected before sending auth ack",
repr(initiator.remote))
ephemeral_pubkey, responder_nonce = initiator.decode_auth_ack_message(
auth_ack)
aes_secret, mac_secret, egress_mac, ingress_mac = initiator.derive_secrets(
initiator_nonce, responder_nonce, ephemeral_pubkey, auth_init,
auth_ack)
return aes_secret, mac_secret, egress_mac, ingress_mac
async def connect(self, aio_loop):
self._allow_states(*ServiceState.halted_states())
self.script = Queue(loop=aio_loop)
self.in_buffer = StreamReader(loop=aio_loop)
self.out_buffer = StreamReader(loop=aio_loop)
# Transfer deferred actions.
for action in self._deferred_actions:
self._add_action_now(action)
self._deferred_actions = []
# If the client wishes to delay startup of the scripted actions (e.g. to register an EventStream first),
# we set this up here.
if self.defer_script:
self._defer_event = Event(loop=aio_loop)
# We're inside of an asyncio task, but need a synchronous start. Hack up
# an asyncio version of it.
started = Event(loop=aio_loop)
def async_start():
self.sync_start()
self.aio_loop.call_soon_threadsafe(started.set)
Thread(target=async_start).start()
await started.wait()
return self.out_buffer, self
async def tcp_recv(reader: asyncio.StreamReader,
delimiter: bytes = b'\n',
timeout=None) -> str:
"""
Receives string result. Handle Incomplete error on your own.
"""
try:
data_length: bytes = await asyncio.wait_for(
reader.readuntil(delimiter), timeout=timeout)
except TypeError:
msg = "tcp_recv: expects"
if not isinstance(delimiter, bytes):
print(msg,
f"<bytes> for delimiter, got {type(delimiter)} instead.")
if not isinstance(timeout, Number) and timeout is not None:
print(msg,
f"<numbers> for delimiter, got {type(timeout)} instead.")
raise
# why this is ignored??? why I am getting ConnReset directly??
except ConnectionResetError:
print("tcp_recv: Disconnected from Server.")
raise
except asyncio.IncompleteReadError:
print("tcp_recv: Incomplete read error.")
raise
data = await asyncio.wait_for(reader.readexactly(
int(data_length.strip(delimiter))),
timeout=timeout)
return data.decode()
async def _handle_client(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
while not reader.at_eof():
try:
array = await reader.readline()
if reader.at_eof():
break
assert array[:1] == b'*'
count = parse_int(array[1:-2])
items = [await read_bulk_string(reader) for _ in range(count)]
command = items[0]
params = items[1:]
try:
with timeout(REDIS_SERVER_TIMEOUT):
await self._execute(command, params, writer)
except Exception:
logger.exception(f"Command failed: {command}")
writer.write(encode_error("Command failed"))
except Exception:
logger.exception(f"Invalid command: {array}")
writer.write(encode_error("Invalid command"))
writer.close()
async def read_stream(reader: StreamReader,
writer: StreamWriter,
log: bool = True) -> AsyncGenerator[str, None]:
"""Reading line from stream."""
reconnect_counter = 0
try:
while not reader.at_eof():
data = await asyncio.wait_for(reader.readline(), timeout=60)
message = data.decode('utf-8').strip()
if log:
logger.debug(message)
yield message
except asyncio.TimeoutError:
logger.exception(f'TimeoutError -> StreamReader timeout get message',
exc_info=False)
except Exception as ex:
logger.exception(f'OtherError -> {ex.__class__.__name__}: {ex}',
exc_info=False)
finally:
reconnect_counter += 1
await asyncio.sleep(60 if reconnect_counter > 60 else reconnect_counter
)
async def proxy(self, client_reader: StreamReader,
client_writer: StreamWriter, remote_reader: StreamReader,
remote_writer: StreamWriter):
client_read_task = create_task(client_reader.read(READ_BYTES_DEFAULT))
remote_read_task = create_task(remote_reader.read(READ_BYTES_DEFAULT))
while client_read_task and remote_read_task:
done, pending = await asyncio.wait(
{client_read_task, remote_read_task},
return_when=FIRST_COMPLETED)
if client_read_task in done:
client_read_task = await self._proxy_connection(
in_read=client_read_task,
out_read=remote_read_task,
in_reader=client_reader,
out_writer=remote_writer)
if remote_read_task in done:
remote_read_task = await self._proxy_connection(
in_read=remote_read_task,
out_read=client_read_task,
in_reader=remote_reader,
out_writer=client_writer)
if client_read_task:
client_read_task.cancel()
if remote_read_task:
remote_read_task.cancel()
remote_writer.close()
async def handler(reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
nonlocal data
header_size = struct.calcsize('!L')
while not reader.at_eof():
try:
header = await reader.readexactly(header_size)
except asyncio.IncompleteReadError:
break
payload_size = struct.unpack("!L", header)[0]
try:
payload = await reader.readexactly(payload_size)
except asyncio.IncompleteReadError:
break
for metric in pickle.loads(payload):
data.append(metric)
if len(data) == count:
event.set()
writer.close()
reader.feed_eof()
async def json_lines_with_timeout(reader: asyncio.StreamReader, timeout: 'seconds' = DEFAULT_TIMEOUT):
while not reader.at_eof():
line = await asyncio.wait_for(reader.readline(), timeout)
try:
yield json.loads(line)
except json.JSONDecodeError as e:
pass
async def _connect_streams(reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
queue: "asyncio.Queue[int]",
token: CancelToken) -> None:
try:
while not token.triggered:
if reader.at_eof():
break
try:
size = queue.get_nowait()
except asyncio.QueueEmpty:
await asyncio.sleep(0)
continue
data = await token.cancellable_wait(reader.readexactly(size))
writer.write(data)
queue.task_done()
await token.cancellable_wait(writer.drain())
except OperationCancelled:
pass
finally:
writer.write_eof()
if reader.at_eof():
reader.feed_eof()
async def decrypt_stream(self, reader: StreamReader) -> StreamReader:
data_length = await reader.read(ENCRYPTED_DATA_LENGTH)
if not data_length:
raise SecurityError('Connection closed')
data_length_int = int.from_bytes(data_length,
byteorder='little') + AUTH_TAG_LENGTH
encrypted_data = await reader.read(data_length_int)
chacha = ChaCha20Poly1305(self.context['decrypt_key'])
nonce = b'\x00\x00\x00\x00' + self.encrypted_request_count.to_bytes(
8, byteorder='little')
try:
decrypted_data = chacha.decrypt(nonce, encrypted_data, data_length)
except InvalidTag:
decrypted_data = None
if not decrypted_data:
raise SecurityError('Unable to decrypt encrypted data')
self.encrypted_request_count += 1
decrypted_reader = StreamReader()
decrypted_reader.feed_data(decrypted_data)
return decrypted_reader
async def _wait_for_either_until_neither(
self, stdout: asyncio.StreamReader,
stderr: asyncio.StreamReader) -> None:
"""
Wait for a line of data from either stdout or stderr and log this data as received.
When both are EOF then exit.
:returns: Tuple of stdout, stderr
"""
future_out = asyncio.ensure_future(stdout.readline())
future_err = asyncio.ensure_future(stderr.readline())
pending = set([future_out,
future_err]) # type: typing.Set[asyncio.Future]
done = set() # type: typing.Set[asyncio.Future]
while len(pending) > 0:
done, pending = await asyncio.wait(pending)
for future_done in done:
result = future_done.result().strip()
if len(result) > 0:
line = result.decode(errors='replace')
if future_done == future_err:
future_err = asyncio.ensure_future(stderr.readline())
pending.add(future_err)
self._logger.error(line)
else:
future_out = asyncio.ensure_future(stdout.readline())
pending.add(future_out)
self._logger.info(line.strip())
def _handler(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
logger.debug('handle')
while True:
cmd = Command.unpack((yield from reader.readexactly(len(Command))))
method = cmd.key.lower()
if method not in self.implements:
yield from self.response(
writer,
Code.INTERNAL_SERVER_ERROR,
bytes('ERROR : method `%s` not implemented\n' % method.strip())
)
continue
args_length = Types.ulong.unpack((
yield from reader.readexactly(len(Types.ulong))
))
kwargs = msgpack.loads((yield from reader.readexactly(args_length)))
try:
code, result = yield from asyncio.coroutine(getattr(self, method))(**kwargs)
yield from self.response(writer, code, result)
except Exception as e:
yield from self.response(writer, Code.INTERNAL_SERVER_ERROR, bytes('ERROR : %s' % e))
logger.exception("Exception when handling request %r")
def connection_made(self, transport):
if self.kill_active:
return
sock = transport.get_extra_info('socket')
set_tcp_keepalive(sock,
opts=dict(
tcp_keepalive=True,
tcp_keepalive_idle=self.tcp_keepalive_time,
tcp_keepalive_intvl=self.tcp_keepalive_interval,
tcp_keepalive_cnt=self.tcp_keepalive_probes,
))
# https://eklitzke.org/the-caveats-of-tcp-nodelay
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
if hasattr(socket, 'TCP_USER_TIMEOUT'):
logger.debug('Setting TCP_USER_TIMEOUT to %s',
self.tcp_user_timeout_seconds * 1000)
sock.setsockopt(socket.SOL_TCP, socket.TCP_USER_TIMEOUT,
self.tcp_user_timeout_seconds * 1000)
self.transport = transport
self._reader = StreamReader()
self._reader.set_transport(transport)
loop = asyncio.get_event_loop()
loop.create_task(self.receive_commands())
super().connection_made(transport)
async def _capture_subprocess_output(
stream: asyncio.StreamReader, ) -> List[bytes]:
lines = []
while not stream.at_eof():
line = await stream.readline()
if line or not stream.at_eof():
lines.append(line.rstrip(b'\n'))
return lines
async def connection_loop(execute_rpc: Callable[[Any], Any],
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
logger: logging.Logger,
cancel_token: CancelToken) -> None:
# TODO: we should look into using an io.StrinIO here for more efficient
# writing to the end of the string.
raw_request = ''
while True:
request_bytes = b''
try:
request_bytes = await cancel_token.cancellable_wait(reader.readuntil(b'}'))
except asyncio.LimitOverrunError as e:
logger.info("Client request was too long. Erasing buffer and restarting...")
request_bytes = await cancel_token.cancellable_wait(reader.read(e.consumed))
await cancel_token.cancellable_wait(write_error(
writer,
f"reached limit: {e.consumed} bytes, starting with '{request_bytes[:20]!r}'",
))
continue
raw_request += request_bytes.decode()
bad_prefix, raw_request = strip_non_json_prefix(raw_request)
if bad_prefix:
logger.info("Client started request with non json data: %r", bad_prefix)
await cancel_token.cancellable_wait(
write_error(writer, 'Cannot parse json: ' + bad_prefix),
)
try:
request = json.loads(raw_request)
except json.JSONDecodeError:
# invalid json request, keep reading data until a valid json is formed
logger.debug("Invalid JSON, waiting for rest of message: %r", raw_request)
continue
# reset the buffer for the next message
raw_request = ''
if not request:
logger.debug("Client sent empty request")
await cancel_token.cancellable_wait(
write_error(writer, 'Invalid Request: empty'),
)
continue
try:
result = await execute_rpc(request)
except Exception as e:
logger.exception("Unrecognized exception while executing RPC")
await cancel_token.cancellable_wait(
write_error(writer, "unknown failure: " + str(e)),
)
else:
writer.write(result.encode())
await cancel_token.cancellable_wait(writer.drain())
def handle_stream(reader: StreamReader, writer: StreamWriter):
addr = writer.get_extra_info('peername')
print("Connect from %r" % (addr, ))
labels_list = load_labels('labels/labels.txt')
load_graph('model/my_frozen_graph_okyonsei.pb')
model = load_model()
count = 0
detected_count = 0
last_data_dict = dict()
try:
while True:
count += 1
data = yield from reader.readexactly(8)
user_id, body_len = struct.unpack("!II", data)
data = yield from reader.readexactly(body_len)
last_data = last_data_dict.get(user_id)
last_data_dict[user_id] = data
if last_data is None:
continue
print(count)
wav_data, np_data = convert_pcm_to_wav(last_data + data)
with open(f'data/{count:06}.wav', 'wb') as fout:
fout.write(wav_data)
detected, prediction = run_graph(wav_data, labels_list, 3)
if detected:
if detected_count:
detected_count = 0
continue
try:
start_p = max(np_data.argmax() - 300, 0)
name = find_who(model, np_data[start_p:])
msg = f'안녕하세요. {name} [정확도: {int(prediction * 100)} %]'
msg_encoded = str.encode(msg)
header = struct.pack("!II", user_id, len(msg_encoded))
writer.write(header + msg_encoded)
yield from writer.drain()
detected_count += 1
except Exception:
pass
else:
detected_count = 0
except Exception as e:
print(e)
writer.close()
async def read_lines_until(reader: StreamReader, timeout: float):
lines = []
while not reader.at_eof():
try:
line = await shield(wait_for(reader.readuntil(b'\r'), timeout))
lines.append(line.decode("ascii"))
# pylint: disable=protected-access
except concurrent.futures._base.TimeoutError:
break
return lines
async def test_BufferedReader_overflow():
stream_data = b"1234"
stream_reader = StreamReader()
stream_reader.feed_data(stream_data)
reader = BufferedReader(stream_reader, 1)
result = await reader.read_bytes(4)
assert result == stream_data
async def receive_message(reader: asyncio.StreamReader):
""" Get the initial command message from the agent.
Used only when both containers are not on a shared-kernel """
buf = bytearray()
while len(buf) != 4 and not reader.at_eof():
buf += await reader.read(4 - len(buf))
length = struct.unpack('!I', bytes(buf))[0]
buf = bytearray()
while len(buf) != length and not reader.at_eof():
buf += await reader.read(length - len(buf))
return msgpack.unpackb(bytes(buf), use_list=False)
async def _client_accept(
self,
reader: StreamReader,
writer: StreamWriter,
read_ahead: bytes = None,
) -> None:
""" Accept new clients and inform the tunnel about connections """
host, port = writer.get_extra_info("peername")[:2]
ip = ipaddress.ip_address(host)
# Block connections using the networks
if self.block(ip):
reader.feed_eof()
writer.close()
await writer.wait_closed()
_logger.info("Connection from %s blocked", ip)
return
self.connections[ip].hits += 1
# Create the client object and generate an unique token
client = Connection(reader, writer, self.protocol,
utils.generate_token())
self.add(client)
_logger.info("Client %s connected on %s:%s", client.uuid, host, port)
# Inform the tunnel about the new client
pkg = package.ClientInitPackage(ip, port, client.token)
await self.tunnel.tun_write(pkg)
# Send the buffer read ahead of initialization through the tunnel
if read_ahead:
await self.tunnel.tun_data(client.token, read_ahead)
# Serve data from the client
while True:
data = await client.read(self.chunk_size)
# Client disconnected. Inform the tunnel
if not data:
break
await self.tunnel.tun_data(client.token, data)
if self.server and self.server.is_serving():
pkg = package.ClientClosePackage(client.token)
await self.tunnel.tun_write(pkg)
await self._disconnect_client(client.token)
async def receive_handshake(
self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
) -> None:
ip, socket, *_ = writer.get_extra_info("peername")
remote_address = Address(ip, socket)
if self.peer_pool.chk_dialin_blacklist(remote_address):
Logger.info_every_n(
"{} has been blacklisted, refusing connection".format(remote_address),
100,
)
reader.feed_eof()
writer.close()
expected_exceptions = (
TimeoutError,
PeerConnectionLost,
HandshakeFailure,
asyncio.IncompleteReadError,
HandshakeDisconnectedFailure,
)
try:
await self._receive_handshake(reader, writer)
except expected_exceptions as e:
self.logger.debug("Could not complete handshake: %s", e)
Logger.error_every_n("Could not complete handshake: {}".format(e), 100)
reader.feed_eof()
writer.close()
except OperationCancelled:
self.logger.error("OperationCancelled")
reader.feed_eof()
writer.close()
except Exception as e:
self.logger.exception("Unexpected error handling handshake")
reader.feed_eof()
writer.close()
async def _log_subprocess_output(
pid: int,
stream: asyncio.StreamReader,
logger: logging.Logger,
level: int,
log_processor: Optional[Callable[[str], Tuple[str, int]]] = None,
) -> List[bytes]:
while not stream.at_eof():
line = await stream.readline()
if line or not stream.at_eof():
log_line = line.rstrip(b'\n').decode()
if log_processor is not None:
log_line, level = log_processor(log_line)
logger.log(level, log_line, extra={"process": pid})
return []
async def _wait_for_event(
self,
event_name: str,
event_type: str,
reader: asyncio.StreamReader,
) -> Any:
matching = self._match_event(
event_name,
event_type,
)
if matching:
return matching
while True:
data = await reader.readline()
line = data.decode('utf-8').rstrip()
if not len(line) and reader.at_eof():
raise Exception(
'Reached end of output waiting for {0}/{1}'.format(
event_name,
event_type,
))
log.info(line)
event = json.loads(line)
matching = self._match_event(
event_name,
event_type,
json.loads(line),
)
if matching:
return matching
async def _pipe(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
try:
while not reader.at_eof():
writer.write(await reader.read(self.CHUNK_SIZE))
finally:
writer.close()
async def handler(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
while not reader.at_eof():
self.data += await reader.read(1)
if self.data:
self.event.set()
async def _get_message(cls, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> str:
"""Gets the next message.
Args:
reader: The reader to receive messages from the client.
writer: The writer to send messages to the client.
Returns:
The next message received from the client.
"""
message = ""
try:
line = await asyncio.wait_for(reader.readline(), timeout=5)
message = line.decode("latin-1").strip()
except asyncio.TimeoutError:
pass
# Check that we're still connected if we didn't get a message
if not message and not await cls._send(writer, ["S,SERVER CONNECTED"]):
raise BrokenPipeError("Client disconnected")
return message
async def test_authenticate(transport, protocol):
conn = AsyncTelnetTransport(
"localhost",
auth_password="******",
auth_username="******",
timeout_ops=10,
timeout_socket=5,
)
conn.stdout = StreamReader()
conn.stdin = StreamWriter(transport,
protocol,
reader=None,
loop=asyncio.get_event_loop())
# feed data in w/ the username prompt already in it -- auth should then send the username and
# clear the buffer... then it should read and see the password that we are inserting later
data = b"\xff\xfd\x18\xff\xfd \xff\xfd#username:\xff\xfd'\xff\xfb\x03\xff\xfd\x01\xff\xfd\x1f\xff\xfb\x05\xff\xfd!\xff\xfb\x03\xff\xfb\x01"
conn.stdout.feed_data(data)
async def _put_login_data_in_reader():
# this sleep is enough that we should *always* see a an additional return char that we send
# when we dont get any output on telnet connections
await asyncio.sleep(1.5)
conn.stdout.feed_data(b"Password:"******"lookforthisusername" in conn.stdin.transport.out_buf
assert b"lookforthispassword" in conn.stdin.transport.out_buf
# a return for the username, one for password and one for the extra return we send when we dont
# see the password prompt soon enough
assert conn.stdin.transport.out_buf.count(b"\n") >= 3
async def _listen_for_messages(self, username: str,
reader: StreamReader): #D
try:
while (data := await asyncio.wait_for(reader.readline(),
60)) != b'':
await self._notify_all(f'{username}: {data.decode()}')
await self._notify_all(f'{username} has left the chat\n')
async def test_isalive_false_not_auth():
conn = AsyncTelnetTransport("localhost")
conn._isauthenticated = False
conn.stdout = StreamReader()
data = b"somethingtoread"
conn.stdout.feed_data(data)
assert conn.isalive() is False
async def test_read_timeout():
conn = AsyncTelnetTransport("localhost")
conn.stdout = StreamReader()
conn._stdout_binary_transmission = True
conn.set_timeout(0.01)
with pytest.raises(ScrapliTimeout):
await conn.read()
def __init__(self, headers, parser, transport, **kw):
self.headers = headers
self.parser = parser
self.reader = StreamReader(**kw)
self.reader.set_transport(transport)
self.feed_data = self.reader.feed_data
self.feed_eof = self.reader.feed_eof
async def handle_client(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
addr = writer.get_extra_info('peername')
log.info("Client connected %r", addr)
while not reader.at_eof():
try:
async with async_timeout.timeout(5):
line = await reader.readline()
if line:
metric = line.decode()
name, value, timestamp = metric.split(" ", 3)
timestamp = float(timestamp)
if value == 'nan':
value = None
else:
value = float(value) if '.' in value else int(value)
await self.storage.write_async((name, (timestamp, value)))
except asyncio.CancelledError:
log.info('Client connection closed after timeout')
break
except: # noqa
continue
log.info("Client disconnected %r", addr)
def from_stream(cls, reader: asyncio.StreamReader, fixed_header: MQTTFixedHeader,
variable_header: MQTTVariableHeader):
data = bytearray()
data_length = fixed_header.remaining_length - variable_header.bytes_length
length_read = 0
while length_read < data_length:
buffer = yield from reader.read(data_length - length_read)
data.extend(buffer)
length_read = len(data)
return cls(data)
def test_get(self, loop):
result = 39
result_queue = Queue()
called_queue = Queue()
reader = StreamReader()
unix_socket = UnixSocket(None, loop)
unix_socket.reader = reader
async def readline():
called_queue.put(True)
return unix_socket.encode(result)
reader.readline = readline
async def run():
unix_socket.ready.set()
result = await unix_socket.get()
result_queue.put(result)
loop.run_until_complete(run())
check_queue(called_queue, True)
check_queue(result_queue, result)
def handle_stdin(self, reader: asyncio.StreamReader):
"""
Handle messages from the agent
"""
try:
while not reader.at_eof():
buf = bytearray()
while len(buf) != 4:
buf += yield from reader.read(4 - len(buf))
length = struct.unpack('I', bytes(buf))[0]
buf = bytearray()
while len(buf) != length:
buf += yield from reader.read(length - len(buf))
message = msgpack.unpackb(bytes(buf), encoding="utf8", use_list=False)
yield from self.handle_stdin_message(message)
except asyncio.CancelledError:
return
except KeyboardInterrupt:
return
except:
self._logger.exception("Exception occured while reading stdin")
os._exit(1) # DIE!
def read_fdms_packet(reader: asyncio.StreamReader) -> bytes:
buffer = bytearray()
ba = yield from reader.read(1)
b = ba[0]
if b > 0x7f:
b &= 0x7f
buffer.append(b)
if b == STX:
got_etx = False
while True:
ba = yield from reader.read(1)
b = ba[0]
if b > 0x7f:
b &= 0x7f
buffer.append(b)
if got_etx:
break
else:
got_etx = (b == ETX)
return buffer
def QueueWaiter(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, queue_name: str):
"""
A coroutine for waiting upon new items to be placed into the Queue.
"""
client = writer.get_extra_info("peername")
sclient = ':'.join(str(_) for _ in client)
while True:
try:
data = yield from reader.read(65536)
except ConnectionResetError:
rlogger.info("Client {} closed connection".format(sclient))
return
if not data:
rlogger.info("Client {} closed connection".format(sclient))
return
# Unpack
try:
sub_data = msgpack.unpackb(data, encoding="utf-8")
except (msgpack.UnpackException, ValueError) as e:
rlogger.error("Recieved non-msgpack push from {}".format(sclient))
continue
rlogger.debug("Recieved data from client {}: {}".format(sclient, sub_data))
assert isinstance(sub_data, dict)
action = sub_data.get("action", -1)
if not action == 0:
rlogger.error("Recieved non-push action on push channel from client {} (action: {})"
.format(sclient, action))
continue
# Get data to place
data = sub_data.get("data", None)
if not data:
rlogger.error("Recieved no data on push channel from client {}".format(sclient))
continue
# Increment and get message number
queues[queue_name][0] += 1
msgnum = queues[queue_name][0]
queue = queues[queue_name][1]
# Put it on the queue
assert isinstance(queue, asyncio.Queue)
yield from queue.put([msgnum, data])
# Respond to the client
response = {"msgnum": msgnum, "status": 0}
rlogger.debug("Sending response with message number {}".format(msgnum))
msgpack.pack(response, writer)
def connected_cb(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
"""
A callback for connected clients.
"""
client = writer.get_extra_info("peername")
sclient = ':'.join(str(_) for _ in client)
logger.info("Recieved connection from {}:{}".format(*client))
# Read a subscription message.
try:
sub = yield from reader.read(65536)
except ConnectionResetError:
rlogger.info("Client {} closed connection".format(sclient))
return
if not sub:
logger.error("Client {} terminated connection abnormally".format(sclient))
return
try:
sub_data = msgpack.unpackb(sub)
except (msgpack.UnpackException, ValueError) as e:
logger.error("Recieved unknown subscription message from {}:{}".format(*client))
yield from writer.drain()
writer.close()
return
# Get the data from the subscription message.
if not b'queue' in sub_data:
logger.error("Recieved null queue from {}".format(sclient))
yield from writer.drain()
writer.close()
return
queue_to_sub = sub_data[b"queue"]
action = sub_data.get(b"action", 0)
queue_created = False
if queue_to_sub not in queues:
queues[queue_to_sub] = [0, asyncio.Queue()]
logger.debug("Created queue {}".format(queue_to_sub))
queue_created = True
logger.debug("Client {} subscribed to queue {} in mode {} ({})".format(sclient, queue_to_sub,
action, "push" if not action else "pull"))
if action == 0:
loop.create_task(QueueWaiter(reader, writer, queue_to_sub))
else:
loop.create_task(QueueSender(reader, writer, queue_to_sub))
msgpack.pack({"created": queue_created}, writer)
async def read_lines_matching(reader: StreamReader, *patterns: Pattern):
waiting_for = {}
for (index, pattern) in enumerate(patterns):
if isinstance(pattern, str):
pattern = re.compile(pattern)
waiting_for[pattern] = index
values = [None] * len(waiting_for)
while not reader.at_eof():
if all(values):
break
line = await reader.readuntil(b'\r')
line = line.decode("ascii").strip()
logger.debug("got line: %s", line)
for pattern in waiting_for:
index = waiting_for[pattern]
if not values[index]:
match = pattern.fullmatch(line)
if match:
values[index] = match[1]
return values
class HttpBodyReader():
_expect_sent = None
_waiting = None
def __init__(self, headers, parser, transport, **kw):
self.headers = headers
self.parser = parser
self.reader = StreamReader(**kw)
self.reader.set_transport(transport)
self.feed_data = self.reader.feed_data
self.feed_eof = self.reader.feed_eof
def waiting_expect(self):
'''``True`` when the client is waiting for 100 Continue.
'''
if self._expect_sent is None:
if (not self.reader.at_eof() and
self.headers.has('expect', '100-continue')):
return True
self._expect_sent = ''
return False
def can_continue(self):
if self.waiting_expect():
if self.parser.get_version() < (1, 1):
raise HttpException(status=417)
else:
msg = '%s 100 Continue\r\n\r\n' % http_protocol(self.parser)
self._expect_sent = msg
self.reader._transport.write(msg.encode(DEFAULT_CHARSET))
def fail(self):
if self.waiting_expect():
raise HttpException(status=417)
def read(self, n=-1):
self.can_continue()
return self.reader.read(n=n)
def readline(self):
self.can_continue()
return self.reader.readline()
def readexactly(self, n):
self.can_continue()
return self.reader.readexactly(n)
def fdms_session(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
online = None
''':type: (FdmsHeader, FdmsTransaction)'''
add_on = None
''':type: (FdmsHeader, FdmsTransaction)'''
offline = list()
writer.write(bytes((ENQ,)))
yield from writer.drain()
while True:
# Get Request
attempt = 0
while True:
try:
if attempt > 4:
return
request = yield from asyncio.wait_for(read_fdms_packet(reader), timeout=15.0)
if len(request) == 0:
return
control_byte = request[0]
if control_byte == STX:
lrs = functools.reduce(lambda x, y: x ^ int(y), request[2:-1], int(request[1]))
if lrs != request[-1]:
raise ValueError('LRS sum')
pos, header = parse_header(request)
txn = header.create_txn()
txn.parse(request[pos:-2])
if header.txn_type == FdmsTransactionType.Online.value:
if online is None:
online = (header, txn)
else:
add_on = (header, txn)
else:
offline.append((header, txn))
if header.protocol_type == '2':
break
# Respond with ACK
attempt = 0
writer.write(bytes((ACK,)))
elif control_byte == EOT:
break
# Close session
except asyncio.TimeoutError:
return
# Respond with NAK
except Exception as e:
logging.getLogger(LOG_NAME).debug('Request error: %s', str(e))
attempt += 1
writer.write(bytes((NAK,)))
yield from writer.drain()
if online is None:
return
# Process Transactions & Send Response
for txn in offline:
rs = process_txn(txn)
offline.clear()
if add_on is not None:
process_add_on_txn(online, add_on)
add_on = None
rs = process_txn(online)
# Send Response
rs_bytes = rs.response()
if rs.action_code == FdmsActionCode.HostSpecificPoll or rs.action_code == FdmsActionCode.RevisionInquiry:
writer.write(rs_bytes)
yield from writer.drain()
else:
attempt = 0
while True:
if attempt >= 4:
return
writer.write(rs_bytes)
yield from writer.drain()
control_byte = 0
try:
while True:
rs_head = yield from asyncio.wait_for(reader.read(1), timeout=4.0)
if len(rs_head) == 0:
return
control_byte = rs_head[0] & 0x7f
if control_byte == ACK:
break
elif control_byte == NAK:
break
# Close session
except asyncio.TimeoutError as e:
return
if control_byte == ACK:
break
else:
attempt += 1
if online[0].wcc in {'B', 'C'}:
# Send ENQ
writer.write(bytes((ENQ,)))
yield from writer.drain()
continue
else:
break
writer.write(bytes((EOT,)))
yield from writer.drain()
if writer.can_write_eof():
writer.write_eof()
