async def relay(
dreader: asyncio.StreamReader,
dwriter: asyncio.StreamWriter,
ureader: asyncio.StreamReader,
uwriter: asyncio.StreamWriter,
) -> None:
"""Pass data/EOF from dreader to uwriter, and ureader to dwriter.
Both writers are ensured to be closed upon exiting this function.
"""
_logger.debug(
'Relaying %r <=> %r', dwriter.get_extra_info('peername'),
uwriter.get_extra_info('peername'))
utask = asyncio.create_task(_relay_data_side(dreader, uwriter))
dtask = asyncio.create_task(_relay_data_side(ureader, dwriter))
async with contexts.aclosing_multiple_writers(dwriter, uwriter):
try:
await asyncio.gather(utask, dtask)
_logger.debug(
'Relay %r <=> %r ended normally',
dwriter.get_extra_info('peername'),
uwriter.get_extra_info('peername'))
except:
dtask.cancel()
utask.cancel()
raise
finally:
await asyncio.wait({dtask, utask})
for t in (dtask, utask):
if t.exception():
_logger.debug(
'Relay task %r caught exception %r', t, t.exception())
def test_put(self, loop):
result = 39
result_queue = Queue()
called_write = Queue()
called_write_eof = Queue()
writer = StreamWriter(None, None, None, None)
unix_socket = UnixSocket(None, loop)
unix_socket.writer = writer
def write(data):
called_write.put(True)
if data != b'\n':
result_queue.put(unix_socket.decode(data))
def write_eof():
called_write_eof.put(True)
writer.write = write
writer.write_eof = write_eof
async def run():
unix_socket.ready.set()
await unix_socket.put(result)
loop.run_until_complete(run())
check_queue_multi(called_write, [True] * 2)
check_queue(result_queue, result)
async def handle_client_connection(
adapter: adapters.ClientAdapter,
transport: str,
upstream_host: str,
upstream_port: int,
args: Dict[str, str],
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
handler_logger.debug(
'Accepted connection for transport %s from %r on %r',
transport,
writer.get_extra_info('peername'), writer.get_extra_info('sockname'))
async with contexts.log_unhandled_exc(handler_logger), \
contexts.aclosing_multiple_writers(writer) as writers:
try:
ureader, uwriter = await adapter.open_transport_connection(
transport, upstream_host, upstream_port, args)
except exceptions.PTConnectError as e:
handler_logger.warning(
'PT reported error while connecting to upstream '
'(%r, %r): %r', upstream_host, upstream_port, e)
writer.transport.abort()
return
writers.add(uwriter)
logname = (f'{writer.get_extra_info("peername")!r} ==> '
f'({upstream_host!r}, {upstream_port})')
handler_logger.info('[%s] %s', transport, logname)
try:
await relays.relay(reader, writer, ureader, uwriter)
except OSError as e:
handler_logger.warning(
'[%s] %s caught %r', transport, logname, e)
async def negotiate_socks4_userid(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
host: Union[str, ipaddress.IPv4Address, ipaddress.IPv6Address],
port: int,
args: Optional[Dict[str, str]],
) -> None:
try:
host = ipaddress.IPv4Address(host)
except ValueError:
raise ValueError('SOCKS4 only supports IPv4 address')
if args:
args_bytes = encode_args(args)
else:
args_bytes = b''
writer.write(b''.join((
b'\x04', # ver
enums.SOCKS4Command.CONNECT,
port.to_bytes(2, 'big'),
host.packed,
args_bytes,
b'\0',
)))
buf = await reader.readexactly(8)
assert buf[0] == 0, 'Invalid SOCKS4 reply version'
reply = enums.SOCKS4Reply(buf[1:2])
if reply is not enums.SOCKS4Reply.GRANTED:
raise exceptions.PTSOCKS4ConnectError(reply)
async def __call__(self, reader: StreamReader,
writer: StreamWriter) -> None:
conn = ManageSieveConnection(self._config, reader, writer)
try:
await conn.run(self._login)
finally:
writer.close()
async def send(self, writer: asyncio.StreamWriter):
"""
Future to send a message over the WebSocket
:param writer: StreamWriter used to send the message
:return:
"""
opcode = self.opcode
data = self.data
frame = bytearray()
head = 0b00000000
head |= 0b10000000
head |= opcode
frame.append(head)
next_byte = 0b00000000
if data:
payload_length = len(data)
else:
payload_length = 0
if 65535 >= payload_length >= 126:
next_byte |= 126
extended_bytes = struct.pack("!H", payload_length)
elif payload_length > 65535:
next_byte |= 127
extended_bytes = struct.pack("!Q", payload_length)
else:
next_byte |= payload_length
extended_bytes = None
frame.append(next_byte)
if extended_bytes:
frame.extend(extended_bytes)
if data:
frame.extend(data)
writer.write(frame)
await writer.drain()
async def _handshake(initiator: 'HandshakeInitiator', reader: asyncio.StreamReader,
writer: asyncio.StreamWriter
) -> Tuple[bytes, bytes, PreImage, PreImage]:
"""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 reader.read(ENCRYPTED_AUTH_ACK_LEN)
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
def __call__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
'''Handle a request
Coroutine.'''
_logger.debug('New proxy connection.')
session = self._new_session(reader, writer)
self.event_dispatcher.notify(self.Event.begin_session, session)
is_error = False
try:
yield from session()
except Exception as error:
if not isinstance(error, StopIteration):
error = True
if isinstance(error, (ConnectionAbortedError,
ConnectionResetError)):
# Client using the proxy has closed the connection
_logger.debug('Proxy error', exc_info=True)
else:
_logger.exception('Proxy error')
writer.close()
else:
raise
finally:
self.event_dispatcher.notify(self.Event.end_session, session,
error=is_error)
writer.close()
_logger.debug('Proxy connection closed.')
async def __call__(self, reader: StreamReader,
writer: StreamWriter) -> None:
conn = IMAPConnection(self.commands, self._config, reader, writer)
state = ConnectionState(self._login, self._config)
try:
await conn.run(state)
finally:
writer.close()
async def _write_ext_msg(
writer: asyncio.StreamWriter,
command: bytes,
body: bytes,
) -> None:
assert len(command) == 2
body_len = len(body).to_bytes(2, 'big')
writer.write(command + body_len + body)
await writer.drain()
async def _client_wrapper(self,
reader: StreamReader,
writer: StreamWriter) -> None:
try:
return await self._callback(
reader=reader, writer=writer,
)
except asyncio.CancelledError:
pass
finally:
writer.close()
async def _relay_data_side(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
"""Pass data and EOF from reader to writer."""
while True:
buf = await reader.read(BUF_SIZE)
if not buf: # EOF
break
writer.write(buf)
await writer.drain()
writer.write_eof()
await writer.drain()
async def _ext_or_port_handler(
self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
# This callback function should not close writer when exiting. After
# all, the API consumer may decide to stash reader and writer somewhere
# for use later and return from their supplied callback function early.
async with contexts.log_unhandled_exc(self._logger):
try:
auth_result = await self._authenticator.authenticate(
reader, writer)
except (OSError, asyncio.IncompleteReadError) as e:
self._logger.warning(
'Error during ExtOrPort SafeCookie authentication: %r', e)
return
if not auth_result:
self._logger.warning(
'ExtOrPort SafeCookie authentication failed')
return
transport = host = port = None
while True:
command, body = await self._read_ext_msg(reader)
if command == enums.ExtOrPortCommand.DONE:
break
elif command == enums.ExtOrPortCommand.USERADDR:
host, port = str_utils.parse_hostport(body.decode('ascii'))
host = ipaddress.ip_address(host)
elif command == enums.ExtOrPortCommand.TRANSPORT:
transport = body.decode('ascii')
str_utils.validate_transport_name(transport)
else:
self._logger.info(
'Received unknown ExtOrPort command %r, body %r',
command, body)
connection_info = ExtOrPortClientConnection(transport, host, port)
if self._preconnect_cb is not None:
accept = await self._preconnect_cb(connection_info)
else:
accept = True
if not accept:
await self._write_ext_msg(
writer, enums.ExtOrPortReply.DENY, b'')
writer.write_eof()
return
await self._write_ext_msg(writer, enums.ExtOrPortReply.OKAY, b'')
await self._cb(reader, writer, connection_info)
def QueueSender(reader: asyncio, writer: asyncio.StreamWriter, queue_name: str):
"""
A coroutine for pulling items from the Queue to the streams.
"""
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:
slogger.info("Client {} closed connection".format(sclient))
return
# Unpack data
try:
sub_data = msgpack.unpackb(data, encoding='utf-8')
except (msgpack.UnpackException, ValueError) as e:
slogger.error("Recieved non-msgpack pull from {}".format(sclient))
continue
action = sub_data.get("action", -1)
if not action == 1:
slogger.error("Recieved non-pull action on pull channel from client (action: {})"
.format(sclient, action))
continue
queue = queues[queue_name][1]
assert isinstance(queue, asyncio.Queue)
data = yield from queue.get()
slogger.debug("Packing data {} for queue {}".format(data[1], queue_name))
response = {"status": 0, "data": data[1], "msgnum": data[0]}
msgpack.pack(response, writer)
async def authenticate(
self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> bool:
"""(async) Authenticate a connecting client.
Returns:
True if authentication is successful and False otherwise. The
caller is responsible for closing the connection in case of
failure.
"""
writer.write(enums.ExtOrPortAuthTypes.SAFE_COOKIE
+ enums.ExtOrPortAuthTypes.END_AUTH_TYPES)
client_auth_type = await reader.readexactly(1)
if client_auth_type != enums.ExtOrPortAuthTypes.SAFE_COOKIE:
return False
client_nonce = await reader.readexactly(self.nonce_len)
server_nonce = secrets.token_bytes(self.nonce_len)
server_hash = self.hash(b''.join((
self.server_hash_header, client_nonce, server_nonce)))
writer.write(server_hash + server_nonce)
client_hash = await reader.readexactly(self.hash_len)
result = hmac.compare_digest(client_hash, self.hash(b''.join((
self.client_hash_header, client_nonce, server_nonce))))
writer.write(int(result).to_bytes(1, 'big'))
return result
async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
data: bytes = await reader.read(2) # Read ID.
if data != b"\x30\x05":
writer.write(ERROR_ID)
return
logger.info(f"Got ID packets: {data}.")
auth: bytes = await reader.read(DIGEST_SIZE)
logger.info(f"Got digest: {auth}.")
length: int = struct.unpack("!I", await reader.read(4))[0]
data = b""
while len(data) < length:
newdata = await reader.read(length - len(data))
if len(newdata) == 0:
break
data += newdata
logger.info(f"Got message ength: {length}, data: {data}.")
try:
logger.info(f"Decoded: {data.decode('UTF-8')}")
message: Dict[str, Any] = json.loads(data.decode("UTF-8"))
logger.info(f"Loaded: {message}")
# Any of these will throw a KeyError with broken packets.
message["type"], message["meta"], message["cont"]
except:
logger.exception("hrrm")
writer.write(ERROR_PACK)
return
stomach: hmac.HMAC = hmac.new(AUTHKEY, data, sha512)
if not hmac.compare_digest(stomach.digest(), auth):
writer.write(ERROR_HMAC)
return
logger.info(message)
writer.write(ERROR_OK)
for event in events:
try:
await event(message)
except:
logger.exception("Caught exception inside commloop event handler.")
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)
async def _accept(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
addr = writer.get_extra_info('peername')
peer = Peer(addr[0], addr[1])
client = PeerTCPClient(self._our_peer_id, peer)
try:
info_hash = await client.accept(reader, writer)
if info_hash not in self._torrent_managers:
raise ValueError('Unknown info_hash')
except Exception as e:
client.close()
if isinstance(e, asyncio.CancelledError):
raise
else:
logger.debug("%s wasn't accepted because of %r", peer, e)
else:
self._torrent_managers[info_hash].accept_client(peer, client)
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 handle_client(self, reader: StreamReader, writer: StreamWriter) -> None:
connection = h11.Connection(h11.SERVER)
body = None # type: StreamReader
while True:
data = await reader.read(65536)
connection.receive_data(data)
event = connection.next_event()
if event is h11.NEED_DATA:
continue
elif isinstance(event, h11.Request):
headers = CIMultiDict((key.decode('ascii'), value.decode('iso-8859-1'))
for key, value in event.headers)
peername = writer.get_extra_info('peername')
peercert = writer.get_extra_info('peercert')
parsed = urlparse(event.target, allow_fragments=False)
query = unquote(parsed.query.decode('ascii'))
request = HTTPRequest(
event.http_version.decode('ascii'), event.method.decode('ascii'),
parsed.path.decode('utf-8'), query, headers, body, bool(self.tls_context),
peername, peercert)
elif isinstance(event, h11.Data):
body.feed_data(event.data)
elif isinstance(event, h11.EndOfMessage):
body.feed_eof()
async def send_msg(stream: StreamWriter, data: bytes):
size_bytes = len(data).to_bytes(4, byteorder='big')
stream.writelines([size_bytes, data])
await stream.drain()
def send_line_to_writer(writer: asyncio.StreamWriter, line):
print('->', line)
writer.write(line.encode('utf-8') + b'\r\n')
async def negotiate_socks5_userpass(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
host: Union[str, ipaddress.IPv4Address, ipaddress.IPv6Address],
port: int,
args: Optional[Dict[str, str]],
) -> None:
if args:
args_bytes = encode_args(args)
if len(args_bytes) > 255 * 2:
raise ValueError('Encoded args too long')
username = args_bytes[:255]
password = args_bytes[255:]
if not password:
password = b'\0'
writer.write(b'\x05\x01' # SOCKS5, 1 auth method
+ enums.SOCKS5AuthType.USERNAME_PASSWORD)
buf = await reader.readexactly(2)
assert buf[0] == 5, 'Invalid server SOCKS version'
if buf[1:2] != enums.SOCKS5AuthType.USERNAME_PASSWORD:
raise RuntimeError(
f'PT rejected userpass auth method, returned {buf[1:2]!r}')
writer.write(b''.join((
b'\x01', # userpass sub-negotiation version 1
len(username).to_bytes(1, 'big'),
username,
len(password).to_bytes(1, 'big'),
password,
)))
buf = await reader.readexactly(2)
assert buf[0] == 1, 'Invalid server USERPASS sub-negotiation version'
if buf[1] != 0:
raise RuntimeError(
f'PT rejected username/password, returned {buf[1:2]!r}')
else:
writer.write(b'\x05\x01' # SOCKS5, 1 auth method
+ enums.SOCKS5AuthType.NO_AUTH)
buf = await reader.readexactly(2)
assert buf[0] == 5, 'Invalid server SOCKS version'
if buf[1:2] != enums.SOCKS5AuthType.NO_AUTH:
raise RuntimeError(
f'PT rejected noauth auth method, returned {buf[1:2]!r}')
try:
host = ipaddress.ip_address(host)
except ValueError:
host_type = enums.SOCKS5AddressType.DOMAIN_NAME
host_bytes = host.encode('idna')
host_len = len(host_bytes)
if host_len > 255:
raise ValueError('Hostname too long')
host_bytes = host_len.to_bytes(1, 'big') + host_bytes
else:
if host.version == 6:
host_type = enums.SOCKS5AddressType.IPV6_ADDRESS
else:
host_type = enums.SOCKS5AddressType.IPV4_ADDRESS
host_bytes = host.packed
writer.write(b''.join((
b'\x05', # SOCKS5
enums.SOCKS5Command.CONNECT,
b'\0', # reserved
host_type,
host_bytes,
port.to_bytes(2, 'big'),
)))
# buf = version, reply, reserved, addr_type, 1st byte of address
buf = await reader.readexactly(5)
assert buf[0] == 5, 'Invalid server SOCKS version'
reply = enums.SOCKS5Reply(buf[1:2])
if reply is not enums.SOCKS5Reply.SUCCESS:
raise exceptions.PTSOCKS5ConnectError(reply)
assert buf[2] == 0, 'Invalid RSV field'
bind_addr_type = enums.SOCKS5AddressType(buf[3:4])
if bind_addr_type is enums.SOCKS5AddressType.IPV4_ADDRESS:
# consume remaining address and port in one call to readexactly()
await reader.readexactly(-1 + 4 + 2)
elif bind_addr_type is enums.SOCKS5AddressType.IPV6_ADDRESS:
await reader.readexactly(-1 + 16 + 2)
else:
await reader.readexactly(buf[4] + 2)
async def handle_client(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
self.DATA.append(await reader.readline())
writer.close()
async def _receive_handshake(
self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
) -> None:
msg = await self.wait(
reader.read(ENCRYPTED_AUTH_MSG_LEN), timeout=REPLY_TIMEOUT
)
ip, socket, *_ = writer.get_extra_info("peername")
remote_address = Address(ip, socket)
self.logger.debug("Receiving handshake from %s", remote_address)
got_eip8 = False
try:
ephem_pubkey, initiator_nonce, initiator_pubkey = decode_authentication(
msg, self.privkey
)
except DecryptionError:
# Try to decode as EIP8
got_eip8 = True
msg_size = big_endian_to_int(msg[:2])
remaining_bytes = msg_size - ENCRYPTED_AUTH_MSG_LEN + 2
msg += await self.wait(reader.read(remaining_bytes), timeout=REPLY_TIMEOUT)
try:
ephem_pubkey, initiator_nonce, initiator_pubkey = decode_authentication(
msg, self.privkey
)
except DecryptionError as e:
self.logger.debug("Failed to decrypt handshake: %s", e)
return
initiator_remote = Node(initiator_pubkey, remote_address)
responder = HandshakeResponder(
initiator_remote, self.privkey, got_eip8, self.cancel_token
)
responder_nonce = os.urandom(HASH_LEN)
auth_ack_msg = responder.create_auth_ack_message(responder_nonce)
auth_ack_ciphertext = responder.encrypt_auth_ack_message(auth_ack_msg)
# Use the `writer` to send the reply to the remote
writer.write(auth_ack_ciphertext)
await self.wait(writer.drain())
# Call `HandshakeResponder.derive_shared_secrets()` and use return values to create `Peer`
aes_secret, mac_secret, egress_mac, ingress_mac = responder.derive_secrets(
initiator_nonce=initiator_nonce,
responder_nonce=responder_nonce,
remote_ephemeral_pubkey=ephem_pubkey,
auth_init_ciphertext=msg,
auth_ack_ciphertext=auth_ack_ciphertext,
)
connection = PeerConnection(
reader=reader,
writer=writer,
aes_secret=aes_secret,
mac_secret=mac_secret,
egress_mac=egress_mac,
ingress_mac=ingress_mac,
)
# Create and register peer in peer_pool
peer = self.peer_pool.get_peer_factory().create_peer(
remote=initiator_remote, connection=connection, inbound=True
)
if (
peer.remote in self.peer_pool.connected_nodes
or peer.remote.pubkey in self.peer_pool.dialedout_pubkeys
):
self.logger.debug("already connected or dialed, disconnecting...")
await peer.disconnect(DisconnectReason.already_connected)
return
if self.peer_pool.is_full:
await peer.disconnect(DisconnectReason.too_many_peers)
return
elif NO_SAME_IP and not self.peer_pool.is_valid_connection_candidate(
peer.remote
):
await peer.disconnect(DisconnectReason.useless_peer)
return
total_peers = len(self.peer_pool)
inbound_peer_count = len(
[peer for peer in self.peer_pool.connected_nodes.values() if peer.inbound]
)
if (
total_peers > 1
and inbound_peer_count / total_peers > self.allow_dial_in_ratio
):
# make sure to have at least (1-allow_dial_in_ratio) outbound connections out of total connections
await peer.disconnect(DisconnectReason.too_many_peers)
else:
# We use self.wait() here as a workaround for
# https://github.com/ethereum/py-evm/issues/670.
await self.wait(self.do_handshake(peer))
async def write_error(writer: asyncio.StreamWriter, message: str) -> None:
json_error = json.dumps({'error': message})
writer.write(json_error.encode())
await writer.drain()
def _write_tcp_packet(packet: Packet, writer: aio.StreamWriter) -> None:
packet = packet.to_bytes()
writer.write(struct.pack('>H', len(packet)) + packet)
writer.drain()
async def handle_connection(reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
host, port = writer.get_extra_info('peername')
logger = main_logger.getChild(f"{host}:{port}")
logger.info("Connection opened from %s", writer.get_extra_info('peername'))
writer.write("DNEVNIK-RU-BKEND-62-02\n".encode())
writer.write("Посчитайте средний балл для каждого ученика\n".encode())
writer.write(
"Если ученик не набрал оценок за период следует вывести н/а\n".encode(
))
task = Task()
remaining = 450
while remaining and not writer.is_closing():
logger.info(f"{remaining} tasks left")
writer.write(f"Осталось {remaining} заданий\n".encode())
try:
task_s = task.get_task()
writer.write(task_s.encode())
line = await reader.readuntil()
logger.info(f"{len(line)} bytes received")
try:
correct = task.check_task(line)
except Exception as e:
writer.write("Неверный формат ответа\n".encode())
logger.info("presentation error")
continue
if correct:
remaining -= 1
writer.write("Верно\n".encode())
logger.info("correct")
else:
writer.write("Неверно\n".encode())
logger.info("incorrect")
break
except Exception as e:
writer.write("Непредвиденная ошибка\n".encode())
logger.error(e)
break
if remaining <= 0:
logger.info("solved")
flag = FLAG
writer.write(f"Ваш флаг: {flag}\n".encode())
writer.write(f"До свидания!\n".encode())
writer.write_eof()
await writer.wait_closed()
async def close(self, reader: StreamReader, writer: StreamWriter) -> None:
""" Close a StreamReader and StreamWriter """
reader.feed_eof()
writer.close()
await writer.wait_closed()
async def send_protocol(self: object, writer: asyncio.StreamWriter,
message: str, address: tuple) -> None:
""" This coroutine sends protocol to participant. """
writer.write(message)
await writer.drain()
self.logger.info('Sent {} to host: {} at port: {}'.format(message.decode(), *address))
async def to_stream(self, writer: asyncio.StreamWriter):
writer.write(self.to_bytes())
await writer.drain()
self.protocol_ts = datetime.now()
async def _do_process_communication(
self,
chain: Chain,
challenge: bytes32,
initial_form: ClassgroupElement,
ip: str,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
):
disc: int = create_discriminant(challenge,
self.constants.DISCRIMINANT_SIZE_BITS)
try:
# Depending on the flags 'fast_algorithm' and 'sanitizer_mode',
# the timelord tells the vdf_client what to execute.
async with self.lock:
if self.config["fast_algorithm"]:
# Run n-wesolowski (fast) algorithm.
writer.write(b"N")
else:
# Run two-wesolowski (slow) algorithm.
writer.write(b"T")
await writer.drain()
prefix = str(len(str(disc)))
if len(prefix) == 1:
prefix = "00" + prefix
if len(prefix) == 2:
prefix = "0" + prefix
async with self.lock:
writer.write((prefix + str(disc)).encode())
await writer.drain()
# Send (a, b) from 'initial_form'.
for num in [initial_form.a, initial_form.b]:
prefix_l = len(str(num))
prefix_len = len(str(prefix_l))
async with self.lock:
writer.write(
(str(prefix_len) + str(prefix_l) + str(num)).encode())
await writer.drain()
try:
ok = await reader.readexactly(2)
except (asyncio.IncompleteReadError, ConnectionResetError,
Exception) as e:
log.warning(f"{type(e)} {e}")
async with self.lock:
self.vdf_failures.append(chain)
self.vdf_failures_count += 1
return
if ok.decode() != "OK":
return
log.info("Got handshake with VDF client.")
async with self.lock:
self.allows_iters.append(chain)
# Listen to the client until "STOP" is received.
while True:
try:
data = await reader.readexactly(4)
except (
asyncio.IncompleteReadError,
ConnectionResetError,
Exception,
) as e:
log.warning(f"{type(e)} {e}")
async with self.lock:
self.vdf_failures.append(chain)
self.vdf_failures_count += 1
break
msg = ""
try:
msg = data.decode()
except Exception:
pass
if msg == "STOP":
log.info(f"Stopped client running on ip {ip}.")
async with self.lock:
writer.write(b"ACK")
await writer.drain()
break
else:
try:
# This must be a proof, 4 bytes is length prefix
length = int.from_bytes(data, "big")
proof = await reader.readexactly(length)
stdout_bytes_io: io.BytesIO = io.BytesIO(
bytes.fromhex(proof.decode()))
except (
asyncio.IncompleteReadError,
ConnectionResetError,
Exception,
) as e:
log.warning(f"{type(e)} {e}")
async with self.lock:
self.vdf_failures.append(chain)
self.vdf_failures_count += 1
break
iterations_needed = uint64(
int.from_bytes(stdout_bytes_io.read(8),
"big",
signed=True))
y_size_bytes = stdout_bytes_io.read(8)
y_size = uint64(
int.from_bytes(y_size_bytes, "big", signed=True))
y_bytes = stdout_bytes_io.read(y_size)
witness_type = uint8(
int.from_bytes(stdout_bytes_io.read(1),
"big",
signed=True))
proof_bytes: bytes = stdout_bytes_io.read()
# Verifies our own proof just in case
int_size = (self.constants.DISCRIMINANT_SIZE_BITS +
16) >> 4
a = int.from_bytes(y_bytes[:int_size], "big", signed=True)
b = int.from_bytes(y_bytes[int_size:], "big", signed=True)
output = ClassgroupElement(int512(a), int512(b))
time_taken = time.time() - self.chain_start_time[chain]
ips = int(iterations_needed / time_taken * 10) / 10
log.info(
f"Finished PoT chall:{challenge[:10].hex()}.. {iterations_needed}"
f" iters, "
f"Estimated IPS: {ips}, Chain: {chain}")
vdf_info: VDFInfo = VDFInfo(
challenge,
iterations_needed,
output,
)
vdf_proof: VDFProof = VDFProof(
witness_type,
proof_bytes,
)
if not vdf_proof.is_valid(self.constants, initial_form,
vdf_info):
log.error("Invalid proof of time!")
async with self.lock:
self.proofs_finished.append(
(chain, vdf_info, vdf_proof))
except ConnectionResetError as e:
log.info(f"Connection reset with VDF client {e}")
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
"""Client handler"""
try:
addr = writer.get_extra_info("peername")
print("Incoming connection from: {}:{}".format(addr[0], addr[1]))
encryption = Encryptor()
pub = encryption.create_rsa_pair()
writer.write(pub)
await writer.drain()
aes_key = encryption.rsa_decrypt(await reader.read(1024))
encryption.set_aes(aes_key)
writer.write(encryption.aes_encrypt(b"ok"))
await writer.drain()
db_conn = sqlite3.connect("db.sqlite")
db_cursor = db_conn.cursor()
while True:
req = await reader.read(4096)
if not req:
raise ConnectionResetError
data = encryption.aes_decrypt(req)
if data:
if data[0] == "0".encode()[0]:
type, short_name, *key, level, hash = data.split(b":")
key = b":".join(key)
db_cursor.execute(
"INSERT INTO Files "
"(FileName, Hash, AccessLvl, Key) "
"VALUES (?, ?, ?, ?)",
(short_name.decode(), hash.decode(), level.decode(),
key))
db_conn.commit()
writer.write(encryption.aes_encrypt(b"ok"))
await writer.drain()
elif data[0] == "1".encode()[0]:
type, key, hash = data.split(b":")
qu = """SELECT FileName, Key FROM Files WHERE Hash = ?
AND AccessLvl <= (SELECT AccessLvl From Users WHERE Keys = ?)"""
res = db_cursor.execute(
qu, (hash.decode(), key.decode())).fetchone()
if res:
file_name, file_key = res
resp = b":".join((file_name.encode(), file_key))
writer.write(encryption.aes_encrypt(resp))
await writer.drain()
else:
writer.write(encryption.aes_encrypt(b"err"))
writer.close()
except ConnectionResetError as err:
addr = writer.get_extra_info("peername")
print("Connection closed: {}:{}".format(addr[0], addr[1]))
writer.close()
def to_stream(self, writer: asyncio.StreamWriter):
writer.write(self.to_bytes())
yield from writer.drain()
self.protocol_ts = datetime.now()
def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
self.r = reader
self.w = writer
self.__socket = writer.get_extra_info("socket")
def send_object(obj: Any, writer: asyncio.StreamWriter):
data = pickle.dumps(obj)
length_data = struct.pack(ControlServer.LENGTH_FMT, len(data))
writer.write(length_data)
writer.write(data)
async def _do_process_communication(
self,
chain: Chain,
challenge: bytes32,
initial_form: ClassgroupElement,
ip: str,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
# Data specific only when running in bluebox mode.
bluebox_iteration: Optional[uint64] = None,
header_hash: Optional[bytes32] = None,
height: Optional[uint32] = None,
field_vdf: Optional[uint8] = None,
):
disc: int = create_discriminant(challenge,
self.constants.DISCRIMINANT_SIZE_BITS)
try:
# Depending on the flags 'fast_algorithm' and 'sanitizer_mode',
# the timelord tells the vdf_client what to execute.
async with self.lock:
if self.sanitizer_mode:
writer.write(b"S")
else:
if self.config["fast_algorithm"]:
# Run n-wesolowski (fast) algorithm.
writer.write(b"N")
else:
# Run two-wesolowski (slow) algorithm.
writer.write(b"T")
await writer.drain()
prefix = str(len(str(disc)))
if len(prefix) == 1:
prefix = "00" + prefix
if len(prefix) == 2:
prefix = "0" + prefix
async with self.lock:
writer.write((prefix + str(disc)).encode())
await writer.drain()
# Send initial_form prefixed with its length.
async with self.lock:
writer.write(
bytes([len(initial_form.data)]) + initial_form.data)
await writer.drain()
try:
ok = await reader.readexactly(2)
except (asyncio.IncompleteReadError, ConnectionResetError,
Exception) as e:
log.warning(f"{type(e)} {e}")
async with self.lock:
self.vdf_failures.append(chain)
self.vdf_failures_count += 1
return
if ok.decode() != "OK":
return
log.info("Got handshake with VDF client.")
if not self.sanitizer_mode:
async with self.lock:
self.allows_iters.append(chain)
else:
async with self.lock:
assert chain is Chain.BLUEBOX
assert bluebox_iteration is not None
prefix = str(len(str(bluebox_iteration)))
if len(str(bluebox_iteration)) < 10:
prefix = "0" + prefix
iter_str = prefix + str(bluebox_iteration)
writer.write(iter_str.encode())
await writer.drain()
# Listen to the client until "STOP" is received.
while True:
try:
data = await reader.readexactly(4)
except (
asyncio.IncompleteReadError,
ConnectionResetError,
Exception,
) as e:
log.warning(f"{type(e)} {e}")
async with self.lock:
self.vdf_failures.append(chain)
self.vdf_failures_count += 1
break
msg = ""
try:
msg = data.decode()
except Exception:
pass
if msg == "STOP":
log.info(f"Stopped client running on ip {ip}.")
async with self.lock:
writer.write(b"ACK")
await writer.drain()
break
else:
try:
# This must be a proof, 4 bytes is length prefix
length = int.from_bytes(data, "big")
proof = await reader.readexactly(length)
stdout_bytes_io: io.BytesIO = io.BytesIO(
bytes.fromhex(proof.decode()))
except (
asyncio.IncompleteReadError,
ConnectionResetError,
Exception,
) as e:
log.warning(f"{type(e)} {e}")
async with self.lock:
self.vdf_failures.append(chain)
self.vdf_failures_count += 1
break
iterations_needed = uint64(
int.from_bytes(stdout_bytes_io.read(8),
"big",
signed=True))
y_size_bytes = stdout_bytes_io.read(8)
y_size = uint64(
int.from_bytes(y_size_bytes, "big", signed=True))
y_bytes = stdout_bytes_io.read(y_size)
witness_type = uint8(
int.from_bytes(stdout_bytes_io.read(1),
"big",
signed=True))
proof_bytes: bytes = stdout_bytes_io.read()
# Verifies our own proof just in case
form_size = ClassgroupElement.get_size(self.constants)
output = ClassgroupElement.from_bytes(y_bytes[:form_size])
if not self.sanitizer_mode:
time_taken = time.time() - self.chain_start_time[chain]
ips = int(iterations_needed / time_taken * 10) / 10
log.info(
f"Finished PoT chall:{challenge[:10].hex()}.. {iterations_needed}"
f" iters, "
f"Estimated IPS: {ips}, Chain: {chain}")
vdf_info: VDFInfo = VDFInfo(
challenge,
iterations_needed,
output,
)
vdf_proof: VDFProof = VDFProof(
witness_type,
proof_bytes,
self.sanitizer_mode,
)
if not vdf_proof.is_valid(self.constants, initial_form,
vdf_info):
log.error("Invalid proof of time!")
if not self.sanitizer_mode:
async with self.lock:
self.proofs_finished.append(
(chain, vdf_info, vdf_proof))
else:
async with self.lock:
writer.write(b"010")
await writer.drain()
assert header_hash is not None
assert field_vdf is not None
assert height is not None
response = timelord_protocol.RespondCompactProofOfTime(
vdf_info, vdf_proof, header_hash, height,
field_vdf)
if self.server is not None:
message = make_msg(
ProtocolMessageTypes.
respond_compact_vdf_timelord, response)
await self.server.send_to_all([message],
NodeType.FULL_NODE)
except ConnectionResetError as e:
log.info(f"Connection reset with VDF client {e}")
async def send_to_stream(msg_q: asyncio.Queue,
stream_out: asyncio.StreamWriter):
msg = await msg_q.get()
stream_out.write(msg)
await stream_out.drain()
async def handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
no = 0
IS_IMG = False
client = {'reader': reader, 'writer': writer}
clients.append(client)
print('Con One')
StringImg = b''
Model = ''
Date = ''
while True:
if IS_IMG == True:
while True:
try:
data: bytes = await reader.read(1024)
if not data:
break
StringImg += data
if b'<eof>' in StringImg:
IS_IMG = False
for i in clients:
if i == client:
continue
i['writer'].write(StringImg)
await i['writer'].drain()
StringImg = b''
break
#StringImg = StringImg.replace(b'<eof>',b'')
#StringImg = StringImg.replace(b'<sof>',b'')
#저장용으로 바꿀 공간=
except:
print('image error')
else:
data: bytes = await reader.read(1024)
if data != None:
if b'<sof>' in data:
StringImg += data
IS_IMG = True
data = b''
continue
peername = writer.get_extra_info('peername')
print(f"[S] received: {len(data)} bytes from {peername}")
cli_mes = data.decode('UTF-8') #client에서 보내준 data를 decode
print(cli_mes)
if cli_mes.startswith('TOC'): # Rasp에서 C#으로 보내줄 data 서버에 출력
print("[Rasp_Client] message: {}".format(cli_mes[3:]))
elif cli_mes.startswith('TOR'): # C#에서 Rasp로 보내줄 data 서버에 출력
print("[C#_Client] message: {}".format(cli_mes[3:]))
no += 1
#해당 index에 값을 넣어주기 위해 추가
conn = sqlite.connect("DB.db") # DB파일에 연결
c = conn.cursor() # 커서획득
if cli_mes.startswith('TEMPS'):
sendf = cli_mes.encode('UTF-8')
writer.write(sendf)
await writer.drain()
if cli_mes.startswith('TORDatesStand'):
payload = cli_mes.encode('UTF-8')
for i in clients:
await asyncio.sleep(0.3)
if i == client:
continue
print(payload)
i['writer'].write(payload)
await i['writer'].drain()
print("[Server] sent: {}".format(cli_mes))
insert_cli_mes = cli_mes.replace('TORDatesStand,',
'').split(',')
index_cli_mes = insert_cli_mes[1]
Model = insert_cli_mes[0]
Date = insert_cli_mes[4]
print(insert_cli_mes)
for no in range(1,
int(index_cli_mes) +
1): #로트 크기만큼 Unit_date 값 넣어주기
c.execute(
'''INSERT INTO Unit_factory(Model, Unit_no, Unit_Date)
VALUES('%s', %s, %s)''' %
(Model, no, Date))
c.execute(
'''INSERT INTO Result(Model, hStandard, vStandard, ResultDate)
VALUES('%s', %s, %s, %s)''' %
(insert_cli_mes[0], insert_cli_mes[2],
insert_cli_mes[3], Date))
elif cli_mes.startswith('TOCUnit_no'):
insert_cli_mes = cli_mes.replace('TOC', '').split(',')
for k in insert_cli_mes:
k = 'TOC' + k
payload = k.encode('UTF-8')
for i in clients:
i['writer'].write(payload)
await i['writer'].drain()
time.sleep(0.1)
print("[Server] sent: {}".format(cli_mes))
print(insert_cli_mes)
#DB파일에 저장
c.execute('''UPDATE Unit_factory SET
Unit_horizon = %s, Unit_vertical = %s,
Unit_hpass = %s, Unit_vpass = %s,
Unit_temp = %s
WHERE Unit_no = %s and Unit_date = %s''' %
(insert_cli_mes[1][12:], insert_cli_mes[2][13:],
insert_cli_mes[3][10:], insert_cli_mes[4][10:],
insert_cli_mes[5][4:], insert_cli_mes[0][7:],
insert_cli_mes[6][9:]))
elif cli_mes.startswith('TOCAQL_hpass'):
print("[Server] sent1: {}".format(cli_mes))
thenTOC = cli_mes.replace('TOC', 'Last')
print(thenTOC)
ttl = thenTOC.encode('UTF-8')
for i in clients:
if i == client:
continue
i['writer'].write(ttl)
await i['writer'].drain()
insert_cli_mes = cli_mes.replace('TOC', '').split(',')
time.sleep(0.3)
print("[Server] sent: {}".format(cli_mes))
try:
c.execute(
'''UPDATE Result SET
AQL_hpass = %s, AQL_vpass = %s,
hSigma = %s, vSigma = %s,
hMean = %s, vMean = %s,
hCp = %s, vCp = %s,
hunpassCount = %s, vunpassCount = %s,
hDefectrate = %s, vDefectrate = %s,
Hadjust = '%s', Vadjust = '%s',
TotalunpassCount = %s, TotalDefectrate = %s,
Result_dateGap= '%s', Result_LOT = %s
WHERE Resultdate = %s and Model = "%s"''' %
(insert_cli_mes[0][9:], insert_cli_mes[1][9:],
insert_cli_mes[2][6:], insert_cli_mes[3][6:],
insert_cli_mes[4][5:], insert_cli_mes[5][5:],
insert_cli_mes[6][3:], insert_cli_mes[7][3:],
insert_cli_mes[8][12:], insert_cli_mes[9][12:],
insert_cli_mes[10][11:], insert_cli_mes[11][11:],
insert_cli_mes[12][7:], insert_cli_mes[13][7:],
insert_cli_mes[14][16:], insert_cli_mes[15][15:],
insert_cli_mes[17][7:], insert_cli_mes[18][3:],
insert_cli_mes[16][4:], insert_cli_mes[19][5:]))
except Exception as e:
print('update exception' + e)
time.sleep(4)
elif cli_mes.startswith('start'):
payload = cli_mes.encode('UTF-8')
for i in clients:
i['writer'].write(payload)
await writer.drain()
print("[Server] sent: {}".format(cli_mes))
conn.commit() # 트랜젝션의 내용을 DB에 반영함
conn.close() # 커서닫기
await asyncio.sleep(1.5) #대기
serv_mes = 'Success'
for i in clients:
i['writer'].write(serv_mes.encode('UTF-8'))
await writer.drain()
#writer.write(serv_mes.encode('UTF-8'))
#await writer.drain()
print("[Server] sent: {}".format(serv_mes))
IS_IMG = False
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,
"reached limit: %d bytes, starting with '%s'" % (
e.consumed,
request_bytes[:20],
),
))
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, from_ipc=True)
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 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()
async def on_connect(_: StreamReader, writer: StreamWriter) -> None:
writer.write(data)
await writer.drain()
async def to_stream(self, writer: asyncio.StreamWriter):
writer.write(self.to_bytes())
await writer.drain()
async def connection_handler(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
# start authentication
_ip = writer.transport.get_extra_info("peername")
_id = self.random_string(16)
self.login_logger.info(f"New connection from {_ip}")
try:
# receive the api key saved in the node and compares it to the local one
proposed_api_key = (await reader.read(1024)).decode()[2:]
except UnicodeDecodeError:
self.login_logger.info(
f"Connection declined for {_ip}. Reason: InvalidEncoding")
writer.close()
return
if proposed_api_key != self.key:
writer.close()
self.login_logger.info(
f"Connection declined for {_ip}. Reason: InvalidApiKey")
return
writer.write(
b"A_ACCEPT"
) # this notifies the client, that it can transfer the config
await writer.drain()
proposed_configuration = await reader.read(
4096
) # this equals 4096 bytes of configuration, which should be enough
try:
proposed_configuration = proposed_configuration.decode("UTF-8")
except UnicodeDecodeError:
self.login_logger.info(
f"Connection declined for {_ip}. Reason: InvalidConfEncoding")
writer.close()
return
try:
configuration = json.loads(proposed_configuration)
new_node = Node()
valid_config = new_node.from_dict(configuration)
if valid_config:
self.nodes[_id] = new_node
else:
self.login_logger.info(
f"Connection declined for {_ip}. Reason: InvalidConf")
return
except json.JSONDecodeError:
self.login_logger.info(
f"Connection declined for {_ip}. Reason: InvalidConf")
return
accepted = False
while not accepted:
writer.write(f"BT_{environ.get('BOT_TOKEN', '')}".encode())
await writer.drain()
if (await reader.read(1024)).decode() == "BT_ACCEPT":
accepted = True
is_ready = False
while not is_ready and not self.stopped:
await asyncio.sleep(0.1)
try:
response = await reader.read(1024)
except BrokenPipeError:
self.login_logger.info(
f"Connection lost to {_ip}. Reason: BPE")
break
if response.decode() == "D_READY":
self.nodes[_id].writer = writer
self.nodes[_id].reader = reader
is_ready = True
self.login_logger.info(f"Connection established to {_ip}")
while not self.stopped and is_ready:
try:
response = await reader.read(1024)
except BrokenPipeError:
del self.nodes[_id]
self.login_logger.info(
f"Connection lost to {_ip}. Reason: BPE")
break
if not response:
print("response none")
break
self._handle_response(response=response.decode())
await asyncio.sleep(0.1)
self.login_logger.info(f"Connection lost to {_ip}. Reason: Basic")
writer.close()
async def _receive_handshake(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> None:
msg = await wait_with_token(
reader.read(ENCRYPTED_AUTH_MSG_LEN),
token=self.cancel_token,
timeout=REPLY_TIMEOUT,
)
ip, socket, *_ = writer.get_extra_info("peername")
remote_address = Address(ip, socket)
self.logger.debug("Receiving handshake from %s", remote_address)
try:
ephem_pubkey, initiator_nonce, initiator_pubkey = decode_authentication(
msg, self.privkey)
except DecryptionError:
# Try to decode as EIP8
msg_size = big_endian_to_int(msg[:2])
remaining_bytes = msg_size - ENCRYPTED_AUTH_MSG_LEN + 2
msg += await wait_with_token(
reader.read(remaining_bytes),
token=self.cancel_token,
timeout=REPLY_TIMEOUT,
)
try:
ephem_pubkey, initiator_nonce, initiator_pubkey = decode_authentication(
msg, self.privkey)
except DecryptionError as e:
self.logger.warn("Failed to decrypt handshake", exc_info=True)
return
# Create `HandshakeResponder(remote: kademlia.Node, privkey: datatypes.PrivateKey)` instance
initiator_remote = Node(initiator_pubkey, remote_address)
responder = HandshakeResponder(initiator_remote, self.privkey,
self.cancel_token)
# Call `HandshakeResponder.create_auth_ack_message(nonce: bytes)` to create the reply
responder_nonce = secrets.token_bytes(HASH_LEN)
auth_ack_msg = responder.create_auth_ack_message(nonce=responder_nonce)
auth_ack_ciphertext = responder.encrypt_auth_ack_message(auth_ack_msg)
# Use the `writer` to send the reply to the remote
writer.write(auth_ack_ciphertext)
await writer.drain()
# Call `HandshakeResponder.derive_shared_secrets()` and use return values to create `Peer`
aes_secret, mac_secret, egress_mac, ingress_mac = responder.derive_secrets(
initiator_nonce=initiator_nonce,
responder_nonce=responder_nonce,
remote_ephemeral_pubkey=ephem_pubkey,
auth_init_ciphertext=msg,
auth_ack_ciphertext=auth_ack_ciphertext)
# Create and register peer in peer_pool
peer = self.peer_class(remote=initiator_remote,
privkey=self.privkey,
reader=reader,
writer=writer,
aes_secret=aes_secret,
mac_secret=mac_secret,
egress_mac=egress_mac,
ingress_mac=ingress_mac,
chaindb=self.chaindb,
network_id=self.network_id)
await self.do_handshake(peer)
async def fix_stream_processor(handler: Handler, shutdown_timeout: float,
reader: AsyncIterator[bytes],
writer: StreamWriter,
cancellation_event: asyncio.Event) -> None:
"""Create a processor for a stream of FIX data.
Args:
handler (Handler): The handler.
shutdown_timeout (float): The time to wait before shutting down.
reader (AsyncIterator[bytes]): The stream reader.
writer (StreamWriter): The stream writer.
cancellation_event (asyncio.Event): An event with which to cancel the processing.
Raises:
RuntimeError: If an invalid event was received.
"""
if cancellation_event.is_set():
return
read_queue: "Queue[Event]" = Queue()
write_queue: "Queue[Event]" = Queue()
async def receive() -> Event:
return await read_queue.get()
async def send(evt: Event) -> None:
await write_queue.put(evt)
await read_queue.put({'type': 'connected'})
state = FixState.OK
message: bytes = b''
reader_iter = reader.__aiter__()
# Create initial tasks.
handler_task = asyncio.create_task(handler(send, receive))
read_task: Task[bytes] = asyncio.create_task(reader_iter.__anext__())
write_task: Task[Event] = asyncio.create_task(write_queue.get())
cancellation_task = asyncio.create_task(cancellation_event.wait())
pending: Set[Future] = {
read_task, write_task, handler_task, cancellation_task
}
# Start the task service loop.
while state == FixState.OK and not cancellation_event.is_set():
# Wait for a task to be completed.
completed, pending = await asyncio.wait(
pending, return_when=asyncio.FIRST_COMPLETED)
# Handle the completed tasks. The pending tasks are left to become completed.
for task in completed:
if task == handler_task:
state = FixState.HANDLER_COMPLETED
continue
elif task == cancellation_task:
state = FixState.CANCELLED
continue
elif task == write_task:
# Fetch the event sent by the handler.
event = write_task.result()
if event['type'] == 'fix':
# Send data to the handler and renew the write task.
LOGGER.debug('Sending "%s"', event["message"])
writer.write(event['message'])
await writer.drain()
write_task = asyncio.create_task(write_queue.get())
pending.add(write_task)
elif event['type'] == 'close':
# Close the connection and exit the task service loop.
writer.close()
state = FixState.HANDLER_CLOSED
continue
else:
LOGGER.debug('Invalid event "%s"', event["type"])
raise RuntimeError(f'Invalid event "{event["type"]}"')
elif task == read_task:
try:
message = cast(bytes, task.result())
LOGGER.debug('Received "%s"', message)
# Notify the client and reset the state.
await read_queue.put({'type': 'fix', 'message': message})
message = b''
# Read the field.
read_task = asyncio.create_task(reader_iter.__anext__())
pending.add(read_task)
except StopAsyncIteration:
state = FixState.EOF
continue
else:
raise AssertionError('Invalid task')
# Attempt to shutdown gracefully.
if state == FixState.HANDLER_COMPLETED:
# When the handler task has finished the session if over.
# Calling done will re-raise an exception.
handler_task.done()
await cancel_await(read_task)
await cancel_await(write_task)
writer.close()
else:
# Notify the client of the disconnection.
await read_queue.put({'type': 'disconnect'})
await cancel_await(write_task)
if state != FixState.EOF:
writer.close()
await cancel_await(read_task)
if not handler_task.cancelled():
LOGGER.info('Waiting %s seconds for the handler to complete.',
shutdown_timeout)
try:
await asyncio.wait_for(handler_task, timeout=shutdown_timeout)
except asyncio.TimeoutError:
LOGGER.error('Cancelling the handler')
await cancel_await(
handler_task, lambda: LOGGER.warning(
'The handler task did not complete and has been cancelled'
))
LOGGER.debug('Shutdown complete.')
async def send_message(message: Union[str, bytes],
writer: asyncio.StreamWriter) -> None:
writer.write((message.encode() if isinstance(message, str) else message) +
SEPARATOR_BINARY)
await writer.drain()
async def submit_message(writer: asyncio.StreamWriter, message: str) -> None:
writer.write(f"{clean_string(message)}\n\n".encode())
await writer.drain()
async def data_loop(
self,
ingress: PacketQueue,
egress: PacketQueue,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
protocol: TCPProtocolV1,
):
def out_task():
return asyncio.create_task(egress.get(),
name=f'send-{protocol.remote.id}')
def in_task():
return asyncio.create_task(
reader.readexactly(protocol.num_to_read()),
name=f'recv-{protocol.remote.id}',
)
def retrieve_exception_cb(task: asyncio.Task):
""" retrieve potential result/exception from orphan tasks """
try:
result = task.result()
self.logger.debug(f'orphan task {task!r} done with {result}')
except:
self.logger.debug(f'orphan task {task!r} got exception',
exc_info=True)
outbound = out_task()
inbound = in_task()
closing = False
while not closing:
done, pending = await asyncio.wait(
{outbound, inbound}, return_when=asyncio.FIRST_COMPLETED)
if outbound in done:
try:
to_send = outbound.result()
except asyncio.CancelledError:
closing = True
else:
if to_send.origin.id == 1 and to_send.received_from and to_send.received_from.id == 4:
self.logger.warn(f'send {to_send}')
to_send = protocol.encode(to_send)
writer.write(to_send)
self.logger.debug(f'send {b2a_hex(to_send)!s}')
outbound = out_task()
if inbound in done:
try:
bs = inbound.result()
except asyncio.IncompleteReadError as ex:
self.logger.info('remote closing connection.')
bs = ex.partial
closing = True
self.logger.debug(f'received {b2a_hex(bs)!s}')
while True:
pkt, origin_id, verify = protocol.decode(bs)
bs = b''
if pkt is None and verify:
assert (protocol.num_to_read() >
0), 'cannot consume buffer'
break
elif pkt is None and not verify:
# 接收到数据，但是数据签名验证异常
#
self.closing(origin_id)
break
else:
await ingress.put(pkt)
# in case multiple packets received in one reading,
# parse as many packets as possible
if protocol.num_to_read() > 0:
break
if not closing:
inbound = in_task()
for t in [inbound, outbound]:
if not t.done():
# add callback to eliminate "exception not retrieve" error
t.add_done_callback(retrieve_exception_cb)
async def sendMessage(writer: asyncio.StreamWriter, msg: Dict[str, Any]) -> None:
''' Encode and send a message to the server. '''
data = json.dumps(msg)
print('\n<{!s} sending {!r}>\n'.format(datetime.datetime.now(), data))
writer.write(data.encode('utf-8') + b'\n')
await writer.drain()
async def _receive_handshake(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> None:
msg = await self.wait(reader.read(ENCRYPTED_AUTH_MSG_LEN),
timeout=REPLY_TIMEOUT)
ip, socket, *_ = writer.get_extra_info("peername")
remote_address = Address(ip, socket)
self.logger.debug("Receiving handshake from %s", remote_address)
got_eip8 = False
try:
ephem_pubkey, initiator_nonce, initiator_pubkey = decode_authentication(
msg, self.privkey)
except DecryptionError:
# Try to decode as EIP8
got_eip8 = True
msg_size = big_endian_to_int(msg[:2])
remaining_bytes = msg_size - ENCRYPTED_AUTH_MSG_LEN + 2
msg += await self.wait(reader.read(remaining_bytes),
timeout=REPLY_TIMEOUT)
try:
ephem_pubkey, initiator_nonce, initiator_pubkey = decode_authentication(
msg, self.privkey)
except DecryptionError as e:
self.logger.debug("Failed to decrypt handshake: %s", e)
return
initiator_remote = Node(initiator_pubkey, remote_address)
responder = HandshakeResponder(initiator_remote, self.privkey,
got_eip8, self.cancel_token)
responder_nonce = secrets.token_bytes(HASH_LEN)
auth_ack_msg = responder.create_auth_ack_message(responder_nonce)
auth_ack_ciphertext = responder.encrypt_auth_ack_message(auth_ack_msg)
# Use the `writer` to send the reply to the remote
writer.write(auth_ack_ciphertext)
await self.wait(writer.drain())
# Call `HandshakeResponder.derive_shared_secrets()` and use return values to create `Peer`
aes_secret, mac_secret, egress_mac, ingress_mac = responder.derive_secrets(
initiator_nonce=initiator_nonce,
responder_nonce=responder_nonce,
remote_ephemeral_pubkey=ephem_pubkey,
auth_init_ciphertext=msg,
auth_ack_ciphertext=auth_ack_ciphertext)
# Create and register peer in peer_pool
peer = self.peer_class(
remote=initiator_remote,
privkey=self.privkey,
reader=reader,
writer=writer,
aes_secret=aes_secret,
mac_secret=mac_secret,
egress_mac=egress_mac,
ingress_mac=ingress_mac,
headerdb=self.headerdb,
network_id=self.network_id,
inbound=True,
)
if self.peer_pool.is_full:
peer.disconnect(DisconnectReason.too_many_peers)
elif not self.peer_pool.is_valid_connection_candidate(peer.remote):
peer.disconnect(DisconnectReason.useless_peer)
total_peers = len(self.peer_pool.connected_nodes)
inbound_peer_count = len([
peer for peer in self.peer_pool.connected_nodes.values()
if peer.inbound
])
if total_peers > 1 and inbound_peer_count / total_peers > DIAL_IN_OUT_RATIO:
# make sure to have at least 1/4 outbound connections
peer.disconnect(DisconnectReason.too_many_peers)
else:
# We use self.wait() here as a workaround for
# https://github.com/ethereum/py-evm/issues/670.
await self.wait(self.do_handshake(peer))
def to_stream(self, writer: asyncio.StreamWriter):
writer.write(self.to_bytes())
yield from writer.drain()
async def handleConnection(self, redirectorClientReader: asyncio.StreamReader,
redirectorClientWriter: asyncio.StreamWriter):
""" Handle connections from redirector clients """
line = b''
sport = pport = shost = phost = None
serviceAddress = peerAddress = None
serviceReader = serviceWriter = None
redirectorClientAddress = None
tasks = []
try:
# fix buffering issues (backpressure effect)
redirectorClientWriter.transport.set_write_buffer_limits(0)
# Get client address from socket
peername = redirectorClientWriter.transport.get_extra_info('peername')
redirectorClientAddress = peername[0] + ':' + str(peername[1])
print('Redirector client connected from', redirectorClientAddress, 'to', self.listenAddress)
try:
with async_timeout.timeout(CONNECT_TIMEOUT, loop=self.loop):
try:
# Wait for the 1st packet (line), which is
line = await redirectorClientReader.readline()
ll = len(line)
# Validate message
if ll < 11 or ll > 28 or line[-1:] != b"\n" or b":" not in line:
raise ValueError()
peerData = line[:-1].decode()
phost, pport, sport = peerData.split(':')
# The address of the peer
peerAddress = phost + ':' + pport
pport = int(pport)
# Parse service port
shost, sport = self.servicePorts[int(sport)]
# Replace ports
if sport in self.replacePorts:
sport = self.replacePorts[sport]
serviceAddress = shost + ':' + str(sport)
except ValueError: pass
except KeyError:
error("Error: No service is on the specified port: {!r}".format(line))
raise GoToEnd()
except asyncio.TimeoutError:
error('Error: Client', redirectorClientAddress, 'protocol error, no data in time!')
raise GoToEnd()
except BrokenPipeError:
error('Client has immediately closed connection!')
raise GoToEnd()
if sport is None:
error('Error: Client', redirectorClientAddress, 'protocol error!')
raise GoToEnd()
# Here we know the peer address and port and the destination port as well so we can create the iptables
# rules and connect to
try:
# Create a socket and get a free port to communicate with. We need to use our own socket to be able to
# get listening port and create iptables rules before connection
csock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
csock.bind((shost, 0))
scport = csock.getsockname()[1]
serviceClientAddress = shost + ':' + str(scport)
except socket.error as e:
error("Socket error: ", e)
raise GoToEnd()
try:
with async_timeout.timeout(CONNECT_TIMEOUT, loop=self.loop):
print(peerAddress, "connecting through", redirectorClientAddress + '…')
# Add IPTables rules
res = await Server.Iptables.addNatRules(redirectorClientAddress, serviceClientAddress,
scport, shost, sport, phost, pport)
# Exit if not connected or no rules
if res != 0: raise GoToEnd()
try:
# Use our own socket to connect
csock.setblocking(False)
await self.loop.sock_connect(csock, (shost, sport))
serviceReader, serviceWriter = await asyncio.open_connection(sock=csock)
except (ConnectionError, BrokenPipeError, OSError, GeneratorExit):
error("Error: service connection from", peerAddress, 'to', serviceAddress, 'is failed!')
raise GoToEnd()
except asyncio.TimeoutError:
error('Error: service', serviceAddress, 'connection error timeout!')
raise GoToEnd()
if not serviceReader or not serviceWriter:
error('Error: service', serviceAddress, 'connection error!')
raise GoToEnd()
# Here connection successfull
print(peerAddress, 'through', redirectorClientAddress, 'connected to', serviceAddress)
async def relayStream(reader, writer, otherWriter):
""" Transfer data from reader to writer """
try:
while True:
# Stop if no data has received in time
with async_timeout.timeout(READ_TIMEOUT, loop=self.loop):
try:
await writer.drain()
data = await reader.read(READ_SIZE)
l = len(data)
if l == 0: # EOF
if not otherWriter.transport.is_closing():
if reader == redirectorClientReader:
print('Peer ', peerAddress, 'through', redirectorClientAddress,
'has closed connection to', serviceAddress)
else:
print('Service', serviceAddress,
'has closed connection from peer', peerAddress,
'through', redirectorClientAddress)
break
writer.write(data)
except (ConnectionError, BrokenPipeError):
if reader == redirectorClientReader:
error('Peer', peerAddress, 'through', redirectorClientAddress,
'has disconnected from service', serviceAddress)
else:
error('Service', serviceAddress, 'has disconnected from peer', peerAddress,
'through', redirectorClientAddress)
break
except OSError as err:
error('Error: OS error:', str(err))
except asyncio.TimeoutError:
if reader == redirectorClientReader:
error('Peer', peerAddress, 'through', redirectorClientAddress,
'read timeout occured. Closing connection.')
else:
error('Service', serviceAddress, 'read timeout occured. Closing connection.')
except asyncio.CancelledError: pass
# Close connection
if not writer.transport.is_closing():
await writer.drain()
writer.close()
# To let the socket actually close
await asyncio.sleep(0, loop=self.loop)
# Create relay tasks
tasks = [
asyncio.ensure_future(relayStream(redirectorClientReader, serviceWriter, redirectorClientWriter),
loop=self.loop),
asyncio.ensure_future(relayStream(serviceReader, redirectorClientWriter, serviceWriter),
loop=self.loop)
]
# Stop waiting when any connection endpoint has closed
done, pending = await asyncio.wait(tasks, loop=self.loop, return_when=asyncio.FIRST_COMPLETED)
# Cancel remaining task
for task in pending: task.cancel()
if pending: await asyncio.wait(pending, loop=self.loop, timeout=1)
tasks = []
except asyncio.CancelledError: pass
except GoToEnd: pass
finally:
# If we have pending tasks, close them
if tasks:
for task in tasks: task.cancel()
await asyncio.wait(tasks, loop=self.loop, timeout=1)
# Close connection from the client if still connected
if not redirectorClientWriter.transport.is_closing():
await redirectorClientWriter.drain()
redirectorClientWriter.close()
# Close connection to the service if still connected
if serviceWriter and not serviceWriter.transport.is_closing():
await serviceWriter.drain()
serviceWriter.close()
# Delete iptables rules
if redirectorClientAddress and Server.Iptables.hasNatRule(redirectorClientAddress):
await Server.Iptables.deleteNatRules(redirectorClientAddress)
def accept_fdms_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
asyncio.Task(fdms.fdms_session(reader, writer), loop=loop).add_done_callback(lambda fut: writer.close())
async def handleConnection(self, peerReader: asyncio.StreamReader, peerWriter: asyncio.StreamWriter):
""" Accepts peer connections """
redirectorReader = redirectorWriter = None
tasks = []
try:
# fix buffering issues (backpressure effect)
peerWriter.transport.set_write_buffer_limits(0)
# Get socket info
peername = peerWriter.get_extra_info('peername')
sockname = peerWriter.get_extra_info('sockname')
# Calculate peer and liten address info
peerAddress = peername[0] + ':' + str(peername[1])
listenPort = str(sockname[1])
listenAddress = sockname[0] + ':' + listenPort
print('Peer', peerAddress, 'connected to', listenAddress)
# Connect to redirector server
try:
with async_timeout.timeout(CONNECT_TIMEOUT, loop=self.loop):
try:
redirectorReader, redirectorWriter = await asyncio.open_connection(self.redirectorServerHost,
self.redirectorServerPort,
ssl=False)
# fix buffering issues (backpressure effect)
redirectorWriter.transport.set_write_buffer_limits(0)
# Send peer info to the redirector server as 1st message
redirectorWriter.write(str.encode(peerAddress + ':' + listenPort + '\n'))
except (ConnectionError, BrokenPipeError, GeneratorExit, OSError):
error('Error: Redirecting connection from ' + peerAddress + ' to ' +
listenAddress + ' is failed!')
except asyncio.TimeoutError:
error('Error: Redirecting connection from ' + peerAddress + ' to ' + listenAddress +
' was failed because of timeout!')
# If connection was unsuccessfull
if not redirectorReader or not redirectorWriter: raise GoToEnd()
print("Peer", peerAddress, 'redirected to', self.redirectorServerAddress)
async def relayStream(reader, writer, otherWriter):
""" Transfer data from reader to writer """
try:
while True:
# Stop if no data has received in time
with async_timeout.timeout(READ_TIMEOUT, loop=self.loop):
try:
await writer.drain()
data = await reader.read(READ_SIZE)
l = len(data)
if l == 0: # EOF
if not otherWriter.transport.is_closing():
if reader == peerReader:
print('Peer', peerAddress, 'has closed connection to',
self.redirectorServerAddress)
else:
print('Redirector', self.redirectorServerAddress,
'has closed connection from peer', peerAddress)
break
writer.write(data)
except (ConnectionError, BrokenPipeError):
if reader == peerReader:
error('Peer', peerAddress, 'has disconnected from', self.redirectorServerAddress)
else:
error('Redirector', self.redirectorServerAddress,
'has disconnected from', peerAddress)
break
except OSError as e:
error('Error: OS error:', str(e))
except asyncio.TimeoutError:
if reader == peerReader:
error('Peer', peerAddress, 'read timeout occured. Closing connection.')
else:
error('Redirector', self.redirectorServerAddress, 'read timeout occured. Closing connection.')
except asyncio.CancelledError: pass
# Close connection
if not writer.transport.is_closing():
await writer.drain()
writer.close()
# To let the socket actually close
await asyncio.sleep(0, loop=self.loop)
# Create relay tasks
tasks = [
asyncio.ensure_future(relayStream(peerReader, redirectorWriter, peerWriter), loop=self.loop),
asyncio.ensure_future(relayStream(redirectorReader, peerWriter, redirectorWriter), loop=self.loop)
]
# Stop waiting when any connection endpoint has closed
done, pending = await asyncio.wait(tasks, loop=self.loop, return_when=asyncio.FIRST_COMPLETED)
# Cancel remaining task
for task in pending: task.cancel()
if pending: await asyncio.wait(pending, loop=self.loop, timeout=1)
tasks = []
except BrokenPipeError:
error('Error: Peer has immediately closed connection!')
except asyncio.CancelledError: pass
except GoToEnd: pass
finally:
# If we have pending tasks close them
if tasks:
for task in tasks: task.cancel()
await asyncio.wait(tasks, loop=self.loop, timeout=1)
# Close peer connection if still opened
if not peerWriter.transport.is_closing():
await peerWriter.drain()
peerWriter.close()
# Close redirector connection if still opened
if redirectorWriter and not redirectorWriter.transport.is_closing():
await redirectorWriter.drain()
redirectorWriter.close()
async def _serv_client(
self, intake: asyncio.StreamReader, outlet: asyncio.StreamWriter,
):
loop = asyncio.get_running_loop()
eol = loop.create_future()
try:
addr = outlet.get_extra_info("peername", "<some-peer>")
# prepare the peer object
ident = str(addr)
# outletting currently has no rate limit, maybe add in the future?
# with an unbounded queue, backpressure from remote peer is ignored
# and outgoing packets can pile up locally
poq = asyncio.Queue()
# intaking should create backpressure when handled slowly, so use a
# bounded queue
hoq = asyncio.Queue(maxsize=1)
peer = Peer(ident=ident, eol=eol, posting=poq.put, hosting=hoq.get,)
# per-connection peer module preparation
modu = {"peer": peer}
if self.init is not None:
# call per-connection peer module initialization method
maybe_async = self.init(modu)
if inspect.isawaitable(maybe_async):
await maybe_async
# launch the peer module, it normally forks a concurrent task to
# run a command landing loop
runpy.run_module(
self.service_modu, modu,
)
logger.debug(f"Nedh client peer module {self.service_modu} initialized")
self.clients.publish(peer)
# pump commands in,
# this task is the only one reading the socket
asyncio.create_task(
receivePacketStream(
peer_site=ident, intake=intake, pkt_sink=hoq.put, eos=eol,
)
)
# pump commands out,
# this task is the only one writing the socket
while True:
pkt = await read_stream(eol, poq.get())
if pkt is EndOfStream:
break
await sendPacket(ident, outlet, pkt)
except asyncio.CancelledError:
pass
except Exception as exc:
logger.error("Nedh client caused error.", exc_info=True)
if not eol.done():
eol.set_exception(exc)
finally:
if not eol.done():
eol.set_result(None)
# todo post err (if any) to peer
outlet.close()
await outlet.wait_closed()
async def _handle_request(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
writer.write(self.prompt)
while 1:
writer.write(b'\r\n$amipy> ')
client = writer.get_extra_info('peername')
_c = ':'.join(str(i) for i in client)
try:
await writer.drain()
data = await reader.readline()
msg = data.decode().strip()
if msg == 'quit':
print(
f'*[Server] {time.ctime()} Connection closed at {_c}')
writer.close()
return
elif msg:
resp = self.parse_opt(msg)
print(
f'*[Server] {time.ctime()} Received "{msg}" from {_c}.'
)
writer.write(resp.encode('latin-1'))
except Exception as e:
print(f'*[Server] {time.ctime()} {e} at {_c}')
writer.close()
if not writer.is_closing():
await writer.drain()
else:
writer.close()
return
