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 echo(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
print('Opened a new connection')
try:
while True:
message: bytes = await reader.readline()
message_str: str = message.decode().replace('\r',
'').replace('\n', '')
print(f'Got message: {message_str}')
if message_str in ['quit', '', 'q']:
writer.write(b'Bye bye\n\r')
await writer.drain()
break
writer.write(message.upper())
# Drain нужен для того, чтоб убедиться что запись
# в сокет проведена, т.к. может записаться в буфер
await writer.drain()
except asyncio.CancelledError:
writer.write(
b'Sorry, server is shutting down. Your connection will be closed\r\n'
)
writer.write_eof()
await writer.drain()
print('Task was cancelled')
finally:
print('Leaving connection')
# raise Exception('Test exception')
# Закрывать reader не нужно.
writer.close()
async def _relay_data_side(
self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
log_name: str,
) -> None:
try:
while True:
buf = await reader.read(self._bufsize)
if not buf: # EOF
break
self._logger.debug('%s passing data', log_name)
writer.write(buf)
await writer.drain()
try:
self._logger.debug('%s passing EOF', log_name)
writer.write_eof()
await writer.drain()
except OSError as e:
if e.errno == errno.ENOTCONN:
self._logger.debug(
'%s endpoint already closed when passing EOF',
log_name)
return
raise
except Exception as e:
self._logger.info('%s caught exception: %r', log_name, e)
raise
async def remove_client(self, writer: asyncio.StreamWriter):
''' Close the cient input & output streams '''
if writer.can_write_eof():
writer.write_eof()
writer.close()
await writer.wait_closed()
self._logger.info("Disconnected client")
async def handle_request(
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
) -> None:
firstline = (await reader.readline()).decode().rstrip('\r\n')
if firstline != "evaluation request":
raise RuntimeError(
f"evaluation request mismatch: {firstline!r}")
# print(f"<<< EVALUATION SERVER >>> received header")
request = json.loads(await reader.readline())
# print(f"<<< EVALUATION SERVER >>> received {request}")
seed = request['seed'][0]
params = request['params'][0]
if len(params) != len(space.params):
raise RuntimeError(
f"Parameter mismatch.\n"
f" received: {sorted(params)}\n"
f" expected: {sorted(p.name for p in space.params)}\n")
x = [params[p.name] for p in space.params]
y, _cost = await objective(x, RandomState(seed))
all_x.append(x)
all_y.append(y)
# print(f"<<< EVALUATION SERVER >>> f({x}) = {y}")
writer.write(json.dumps({'y': y}).encode())
writer.write_eof()
async def _server_callback(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> None:
"""Callback when a connection is made to the server
Read the data sent from the client, execute the requested command, and
send the reply back to the client.
"""
try:
logger.debug("Connection made to server")
data = await reader.read()
logger.debug("EOF received by server")
req, is_json = _IPC.unpack(data)
except IPCError:
logger.warning("Invalid data received, closing connection")
else:
rep = self.handler(req)
result = _IPC.pack(rep, is_json=is_json)
logger.debug("Sending result on receive EOF")
writer.write(result)
logger.debug("Closing connection on receive EOF")
writer.write_eof()
finally:
writer.close()
await writer.wait_closed()
async def request_file_handler(self, writer: StreamWriter, file_path: str, block_index: int):
# read data
def read_file():
with open(file_path, mode="r+b") as f:
f.seek(block_index * config.file_block_size)
data = f.read(config.file_block_size)
if config.enable_gzip:
data = zlib.compress(data)
return data
data = await self._loop.run_in_executor(None, read_file)
# encryption
if self._encryption:
salt = get_random_bytes(AES.block_size)
pk = hashlib.scrypt(config.pre_shared_key, salt=salt, n=2 ** 14, r=8, p=1, dklen=32)
cipher_config = AES.new(pk, AES.MODE_GCM)
encrypted, tag = cipher_config.encrypt_and_digest(data)
msg = pickle.dumps({
"salt": salt,
"cipher": encrypted,
"tag": tag,
"nonce": cipher_config.nonce
})
data = msg
writer.write(int.to_bytes(MsgType.RES_FILE.value, 1, "big"))
writer.write(int.to_bytes(len(data), 8, "big"))
writer.write(data)
writer.write_eof()
async def request_index_handler(self, writer: StreamWriter, client_index: dict):
client_ip = writer.get_extra_info("peername")[0]
# response with local index
local_index = pickle.dumps(self._file_mgr.file_index)
writer.write(int.to_bytes(MsgType.RES_INDEX.value, 1, "big"))
writer.write(int.to_bytes(len(local_index), 8, "big"))
writer.write(local_index)
writer.write_eof()
await self.sync(client_index, client_ip)
async def request_index_update_handler(self, writer: StreamWriter, client_index: dict):
client_ip = writer.get_extra_info("peername")[0]
# response
data = b"OK"
writer.write(int.to_bytes(MsgType.RES_INDEX_UPDATE.value, 1, "big"))
writer.write(int.to_bytes(len(data), 8, "big"))
writer.write(data)
writer.write_eof()
await self.sync(client_index, client_ip)
async def handle_client(
self,
reader: asyncio.StreamReader,
writer: asyncio.StreamWriter,
):
self.DATA.append(await reader.readline())
async with self.condition:
self.condition.notify_all()
writer.write_eof()
writer.close()
def run_on_connected(reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
if self.connection_task is not None:
logger.warning(
'Trying to connect, but controller already connected.')
writer.write('Already connected'.encode())
writer.write_eof()
return
self.remote_addr = writer.transport.get_extra_info('peername')
self.connection_task = asyncio.ensure_future(
on_connected(reader, writer))
async def tcp_received(self, reader: aio.StreamReader,
writer: aio.StreamWriter) -> None:
logging.debug(f'Tcp connected {writer.get_extra_info("peername")}')
try:
await aio.wait_for(self._tcp_received(reader, writer),
timeout=self._TCP_TIME_LIMIT)
except Exception:
pass
finally:
writer.write_eof()
writer.close()
logging.debug('Done TCP')
async def _read_fd_write_stream(loop: asyncio.AbstractEventLoop, rfd: int,
writer: StreamWriter) -> None:
f = os.fdopen(rfd, mode="rb")
try:
while True:
b = await loop.run_in_executor(None, f.read, CHUNK_SIZE)
if not b:
break
writer.write(b)
writer.write_eof()
await writer.drain()
finally:
f.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 send_response(
writer: asyncio.StreamWriter, response: Dict[Any, Any],
) -> None:
"""
Sends a response to a connection and then closes writing to that connection
:param writer: StreamWriter to write to
:param response: Dict[Any, Any] response
:return: None
"""
writer.write(bencode.encode(response))
writer.write_eof()
await writer.drain()
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 handle_client_connection(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
service_name = self.__router.determine_service(reader, writer)
responsible_services = list(
filter(lambda s: s.get_service_name() == service_name,
self.__services))
if len(responsible_services) != 0:
for service in responsible_services:
await service.client_connected(reader, writer)
else:
writer.write(ERROR_MESSAGE)
writer.write_eof()
await writer.drain()
writer.close()
async def echo(reader: StreamReader, writer: StreamWriter):
print('New connection.')
try:
while True:
data: bytes = await reader.readline()
if data in [b'', b'quit']:
break
writer.write(data.upper())
await writer.drain()
print('Leaving Connection.')
except CancelledError:
writer.write_eof()
print('Cancelled')
finally:
writer.close()
async def poll_client(reader: StreamReader, process: Process, stdin: StreamWriter):
"""
Continuously waits to receive a stdin message from the client.
"""
while True:
(message_type, body) = await read_client(reader)
if message_type == MESSAGE_STDIN:
if len(body) == 0:
stdin.write_eof()
else:
stdin.write(body)
elif message_type == MESSAGE_SIGNAL:
signal = int.from_bytes(body, "big", signed=False)
print("Got signal", signal)
process.send_signal(signal)
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)
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 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 _reply(self, dwriter: asyncio.StreamWriter, reply: SOCKS5Reply,
host: HostType, port: int) -> None:
if isinstance(host, ipaddress.IPv4Address):
b_addr = SOCKS5AddressType.IPV4_ADDRESS + host.packed
elif isinstance(host, ipaddress.IPv6Address):
b_addr = SOCKS5AddressType.IPV6_ADDRESS + host.packed
else:
b_addr = host.encode('idna')
b_addr = (SOCKS5AddressType.DOMAIN_NAME +
len(b_addr).to_bytes(1, 'big') + b_addr)
dwriter.write(b'\x05' + reply + b'\x00' + b_addr +
port.to_bytes(2, 'big'))
if reply is not SOCKS5Reply.SUCCESS:
dwriter.write_eof()
await dwriter.drain()
self._logger.debug('%r sent reply %s',
dwriter.get_extra_info('peername'), reply)
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 _connection_wrapper(application: core.ApplicationType, client_connections: "Set[asyncio.Task[None]]", container: core.Container, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
"""
Run an ASGI application in an asyncio server.
This function is suitable for passing to ``start_server`` or
``start_unix_server``, with the ``application`` and ``client_connections``
parameters bound via a ``functools.partial`` or similar.
:param application: The ASGI application.
:param client_connections: A set of Task objects, to which this connection
is added on entry to and removed on exit from this function.
:param container: The ASGI container to use.
:param reader: The stream reader for the connection.
:param writer: The stream writer for the connection.
"""
# Add this task to the set of open client connections.
task = asyncio.current_task()
assert task is not None, "_connection_wrapper must be called inside a task"
client_connections.add(task)
try:
try:
# aioscgi.run expects callables to read a chunk and write a chunk,
# the latter taking a drain boolean; adapt the writing side to the
# stream model (the reader is handled with a functools.partial).
async def write_cb(data: bytes, drain: bool) -> None:
writer.write(data)
if drain:
await writer.drain()
# Run the application.
await container.run(application, functools.partial(reader.read, io.DEFAULT_BUFFER_SIZE), write_cb)
except Exception: # pylint: disable=broad-except
logging.getLogger(__name__).error("Uncaught exception in application callable", exc_info=True)
finally:
# Close the connection.
try:
if writer.can_write_eof():
writer.write_eof()
writer.close()
except Exception: # pylint: disable=broad-except
# If something went wrong while closing the connection, there’s
# nothing interesting to report.
pass
finally:
# Remove this task from the set of open client connections.
client_connections.remove(task)
async def login(reader: StreamReader, writer: StreamWriter, name: str) -> bool:
global mysql_pool
if mysql_pool is None:
writer.write(b"database error\n")
return False
bitsize = 256
str_len = bitsize // 4 + 2
client_eph: int
server_eph: int
writer.write(b"exchange\n")
client_eph_str = await reader.readline()
client_eph_str = client_eph_str.strip()
if len(client_eph_str) != str_len:
writer.write(b"too short\n")
return False
try:
client_eph = int(client_eph_str, 16)
except ValueError:
writer.write(b"not a number\n")
return False
server_eph = randbits(bitsize)
server_eph_str = hex(server_eph).encode()
writer.write(server_eph_str + b"\n")
data = hex(server_eph * client_eph).encode()
signature = b64decode(await reader.readline())
async with mysql_pool.acquire() as conn:
async with conn.cursor() as cur:
await cur.execute("select certificate from users where name = %s",
(name, ))
(public_key, ) = await cur.fetchone()
try:
key = crypto.load_certificate(crypto.FILETYPE_PEM, public_key.encode())
crypto.verify(key, signature, data, "sha256")
except crypto.Error as e:
print(f"Verification failed: {e}")
writer.write(b"fail\n")
writer.write_eof()
return False
writer.write(b"ok\n")
return True
async def on_connected(reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
logger.info('Connected from %s', self.remote_addr)
self.begin_control()
while not reader.at_eof():
line = await reader.readline()
line = line.decode().strip()
if not line:
if not reader.at_eof():
logger.warning('Empty command received')
continue
self._process_command(line)
logger.info('Disonnected')
writer.write_eof()
self.end_control()
self.connection_task = None
async def communicate(
self,
reader: StreamReader,
writer: StreamWriter,
):
protocol = Protocol(reader, writer)
request_writer_task = self.create_task(self._request_writer(protocol))
client_host, client_port = writer.get_extra_info("peername")[:2]
if ":" in client_host:
client_host = f"[{client_host}]"
log.info(
"Start communication with tcp://%s:%d",
client_host,
client_port,
)
# noinspection PyBroadException
try:
async for payload in protocol:
if payload.request is not None:
self.create_task(self.execute(protocol, payload.request), )
if payload.response is not None:
self._on_response(payload.response)
except Exception:
log.exception(
"Error when communication tcp://%s:%d/",
client_host,
client_port,
)
finally:
if writer.can_write_eof():
writer.write_eof()
writer.close()
log.info(
"Communication with tcp://%s:%d finished",
client_host,
client_port,
)
await cancel_tasks([request_writer_task])
await writer.wait_closed()
async def _server_callback(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter) -> None:
"""Callback when a connection is made to the server
Read the data sent from the client, execute the requested command, and
send the reply back to the client.
"""
try:
logger.debug("Connection made to server")
data = await reader.read()
logger.debug("EOF received by server")
req, is_json = _IPC.unpack(data)
except IPCError:
logger.warn("Invalid data received, closing connection")
else:
if req[1] == "restart":
# if we are going to restart, close the connection first, as we won't be back
logger.debug("Closing connection on restart")
writer.write_eof()
rep = self.handler(req)
result = _IPC.pack(rep, is_json=is_json)
logger.debug("Sending result on receive EOF")
writer.write(result)
logger.debug("Closing connection on receive EOF")
writer.write_eof()
finally:
# the resoure isn't closed immediately on the close call, but is on
# the next loop iteration, this is exposed as the wait_closed
# method in 3.7, but requires a manual loop iteration in earlier
# versions
writer.close()
if sys.version_info >= (3, 7):
await writer.wait_closed()
else:
await asyncio.sleep(0)
async def handle_req(reader: asyncio.StreamReader,
writer: asyncio.StreamWriter, filename: Path) -> None:
"""Send data."""
try:
LOG.info("Connection started...")
if filename.suffix == ".gz":
LOG.info("Opening gzfile...")
fp_ = gzip.open(filename, "rb")
elif zipfile.is_zipfile(filename):
LOG.info("Opening zipfile...")
zfile = zipfile.ZipFile(filename)
fp_ = zfile.open(zfile.filelist[0], "r")
else:
fp_ = filename.open("rb")
await reader.read(4026)
while True:
line = fp_.readline()
if not line:
LOG.info(
"No remaining line...sending term line and breaking...")
break
writer.write(line)
await writer.drain()
writer.write_eof()
writer.close()
await writer.wait_closed()
LOG.debug("All lines sent...closing...")
except (ConnectionResetError, BrokenPipeError, CancelledError):
LOG.info("Cancel received..shutting down...")
writer.close()
except Exception as err:
LOG.info(f"Unexepcted error! {err}")
try:
writer.close()
except:
pass
LOG.info("Exiting...")
async def ahandle_peer(self, reader: StreamReader, writer: StreamWriter) -> None:
"""Read all DNS queries from the peer stream and schedule their resolution via a DnsResolver instance."""
tasks: Union[List[Task], Set[Task]] = []
wlock = aio.Lock()
logging.debug(f'Got TCP DNS query stream from {writer.transport.get_extra_info("peername")}')
while True:
# Parse a DNS query packet off of the wire
try:
query_size = int.from_bytes(await reader.readexactly(2), 'big')
query = await reader.readexactly(query_size)
# Check if our peer has finished writing to the stream
except aio.IncompleteReadError:
break
# Schedule the processing of the query
tasks.append(aio.create_task(self.ahandle_query(writer, wlock, query)))
# Wait for all scheduled query processing to finish
while tasks:
done, tasks = await aio.wait(tasks, return_when=aio.FIRST_COMPLETED)
for task in done:
error = task.exception()
if error is not None:
logging.warning(f'TCP DNS query resolution encountered an error - {error!r}')
if not writer.is_closing():
# Indicate we are done writing to the stream
if writer.can_write_eof():
writer.write_eof()
# Close the stream
writer.close()
await writer.wait_closed()
async def _handle_msg(self, reader: asyncio.StreamReader,
writer: asyncio.StreamWriter):
"""Handle incoming messages."""
msg = Cmd(await reader.read())
addr = writer.get_extra_info("peername")
print_addr = addr if addr else "localhost"
self.server_logger.debug(f"Received from {print_addr}")
if msg.cmd in self.cmd_to_callback:
self.server_logger.debug(f"Got cmd: {msg.cmd}")
response = await self.cmd_to_callback[msg.cmd](msg)
else:
self.server_logger.warning(
f"Got unrecognized command: {msg.cmd.value}")
response = badResponse()
response.error = ERRORS.UNRECOGNIZED_COMMAND
writer.write(response.get_bytes())
writer.write_eof()
await writer.drain()
self.server_logger.debug(f"Closed connection from {print_addr}")
writer.close()
async def _handle_socks5_connection(self, dreader: asyncio.StreamReader,
dwriter: asyncio.StreamWriter):
log_name = '{!r} <=> ()'.format(
dwriter.transport.get_extra_info('peername'))
try: # catch, log and suppress all exceptions in outermost layer
with ExitStack() as stack:
stack.enter_context(
self._connections.with_this(asyncio.Task.current_task()))
stack.enter_context(finally_close(dwriter))
self._logger.debug('%s accepted downstream connection',
log_name)
# Negotiate incoming SOCKS5 connection
# Authentication
buf = await dreader.readexactly(1) # Version marker
if buf[0] != 5:
raise RuntimeError('%s invalid SOCKS version' % log_name)
buf = await dreader.readexactly(1) # number of auth methods
buf = await dreader.readexactly(buf[0]) # offered auth methods
if SOCKS5AuthType.NO_AUTH not in buf:
self._logger.info('%s did not offer "no auth", offers: %r',
log_name, buf)
dwriter.write(
bytes((5, SOCKS5AuthType.NO_OFFERS_ACCEPTABLE)))
dwriter.write_eof()
await dwriter.drain()
return
dwriter.write(bytes((5, SOCKS5AuthType.NO_AUTH)))
# client command
buf = await dreader.readexactly(4) # ver, cmd, rsv, addr_type
if buf[0] != 5 or buf[2] != 0:
raise RuntimeError('%s malformed SOCKS5 command' %
log_name)
cmd = buf[1]
addr_type = buf[3]
if addr_type == SOCKS5AddressType.IPV4_ADDRESS:
uhost = ipaddress.IPv4Address(
await dreader.readexactly(4)).compressed
elif addr_type == SOCKS5AddressType.IPV6_ADDRESS:
uhost = ipaddress.IPv6Address(
await dreader.readexactly(16)).compressed
elif addr_type == SOCKS5AddressType.DOMAIN_NAME:
buf = await dreader.readexactly(1) # address len
uhost = (await dreader.readexactly(buf[0])).decode('utf-8')
else:
raise RuntimeError('%s illegal address type' % log_name)
uport = int.from_bytes(await dreader.readexactly(2), 'big')
log_name = '{!r} <=> ({!r}, {!r})'.format(
dwriter.transport.get_extra_info('peername'), uhost, uport)
self._logger.debug('%s parsed target address', log_name)
if cmd != SOCKS5Command.CONNECT:
self._logger.info('%s command %r not supported', log_name,
cmd)
dwriter.write(
self._make_socks5_command_reply(
SOCKS5Reply.COMMAND_NOT_SUPPORTED, '0.0.0.0', 0))
dwriter.write_eof()
await dwriter.drain()
return
self._logger.info('%s received CONNECT command', log_name)
# determine detour state of host name
# Since getting and setting detour state is separated quite far,
# multiple connections to the same host will quite possibly
# make the state inconsistent. However that does not have much
# adverse effects on the operation.
detour_state = self._whitelist.state(uhost)
if not detour_state[0]:
self._logger.info('%s try direct connection', log_name)
try:
ureader, uwriter = await asyncio.wait_for(
asyncio.open_connection(
uhost,
uport,
loop=self._loop,
limit=self.RELAY_BUFFER_SIZE),
self.DETOUR_TIMEOUT,
loop=self._loop)
except (OSError, asyncio.TimeoutError) as e:
self._logger.info('%s direct connection error: %r',
log_name, e)
self._whitelist.add_to_temp_wl(uhost)
# continue to making detoured connection
else:
self._logger.info('%s direct connection successful',
log_name)
stack.enter_context(finally_close(uwriter))
dwriter.write(
self._make_socks5_command_reply(
SOCKS5Reply.SUCCESS,
*(uwriter.transport.get_extra_info('sockname')
[:2])))
try:
await self._relay_data(dreader, dwriter, ureader,
uwriter, (uhost, uport))
except UpstreamNetworkError as e:
self._logger.info(
'%s direct connection upstream '
'error during relay: %r', log_name,
e.__cause__)
self._whitelist.add_to_temp_wl(uhost)
raise
self._logger.info(
'%s direct connection completed '
'without error', log_name)
# not going to make a detoured connection and try again
# because retrying means resending all sent data
return
self._logger.info('%s try detoured connection', log_name)
try:
ureader, uwriter = await aiosocks.open_connection(
self._upstream_addr,
self._upstream_auth, (uhost, uport),
remote_resolve=True,
loop=self._loop,
limit=self.RELAY_BUFFER_SIZE)
except (OSError, aiosocks.SocksError) as e:
self._logger.info('%s detour connection error: %r',
log_name, e)
dwriter.write(
self._make_socks5_command_reply(
SOCKS5Reply.GENERAL_SOCKS_SERVER_FAILURE,
'0.0.0.0', 0))
dwriter.write_eof()
await dwriter.drain()
if detour_state[0] == DetourState.TEMP:
self._whitelist.remove_from_temp_wl(detour_state[1])
return
self._logger.info('%s detour connection successful', log_name)
stack.enter_context(finally_close(uwriter))
dwriter.write(
self._make_socks5_command_reply(
SOCKS5Reply.SUCCESS,
*uwriter.transport.get_extra_info('sockname')[:2]))
try:
await self._relay_data(dreader, dwriter, ureader, uwriter,
(uhost, uport))
except UpstreamNetworkError as e:
self._logger.info(
'%s detoured connection upstream error '
'during relay: %r', log_name, e.__cause__)
if detour_state[0] == DetourState.TEMP:
self._whitelist.remove_from_temp_wl(detour_state[1])
raise
self._logger.info(
'%s detoured connection completed without '
'error', log_name)
if detour_state[0] == DetourState.TEMP:
self._whitelist.add_to_perm_wl(detour_state[1])
except asyncio.CancelledError:
self._logger.debug('%s cancelled', log_name)
raise
except (
RuntimeError,
OSError,
UpstreamNetworkError,
) as e: # not logging stack trace for normal errors
self._logger.info('%s %r', log_name, e)
except Exception as e:
self._logger.error('%s %r', log_name, e, exc_info=True)
finally:
self._logger.debug('%s connection done', log_name)
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()
