class ControlledConnection:
def __init__(self, data, limit, leave_open=False):
self._reader = StreamReader(limit)
self._reader.feed_data(data)
if not leave_open:
self._reader.feed_eof()
self._mock_write_data = b""
def get_mock_write_data(self):
return self._mock_write_data
async def read(self, size):
try:
return await self._reader.read(min(size, 100))
except Exception:
raise MalformedDataError
async def readuntil(self, delim):
try:
return await self._reader.readuntil(delim)
except Exception:
raise MalformedDataError
async def readexactly(self, amount):
try:
return await self._reader.readexactly(amount)
except Exception:
raise MalformedDataError
async def write(self, data):
self._mock_write_data += data
def close(self):
pass
def hmland(manager, host, port):
while True:
print("Connecting to hmland")
try:
loop = events.get_event_loop()
reader = StreamReader(limit=_DEFAULT_LIMIT, loop=loop)
protocol = StreamReaderProtocol(reader, loop=loop)
transport, _ = yield from loop.create_connection(lambda: protocol, host, port)
writer = StreamWriter(transport, protocol, reader, loop)
reader, writer = yield from asyncio.open_connection(host, port)
transport._sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
manager.setWriter(writer)
asyncio.async(manager.update_locks())
while True:
line = yield from reader.readline()
if not line:
break
yield from manager.handle(line.decode("UTF-8").strip())
except:
pass
yield from asyncio.sleep(2)
def connection_made(self, transport):
self.reader = StreamReader(loop=self.loop)
self.transport = transport
self.request.remote_addr = self.remoteaddr
self.reader.set_transport(transport)
self.loop.create_task(self.parse())
def gzip_decompress(reader, writer):
new_reader = StreamReader(limit=reader._limit, loop=reader._loop)
task = asyncio.create_task(decompressReader(reader, new_reader))
_feed_eof = new_reader.feed_eof
def new_feed_eof():
if not task.done():
task.cancel()
_feed_eof()
new_reader.feed_eof = new_feed_eof
return new_reader, writer
def make(data, *, limit=None, exc=None):
if limit is not None:
r = StreamReader(loop=event_loop, limit=limit)
else:
r = StreamReader(loop=event_loop)
w = mock_stream_writers()
r.feed_data(data)
if exc is not None:
r.set_exception(exc)
else:
r.feed_eof()
return r, w
def __init__(self, *args, quality_check=False, **kwargs):
"""
Args:
*args:
quality_check (bool): there are times when we only establish a connection to check the quality of the node/address
**kwargs:
"""
self._stream_reader = StreamReader()
self._stream_writer = None
nodemanager = kwargs.pop('nodemanager')
self.client = NeoNode(self, nodemanager, quality_check)
self._loop = events.get_event_loop()
super().__init__(self._stream_reader)
async def handle_connection(self, reader: StreamReader, writer: StreamWriter):
temp_ref = TempRef()
world_packet_manager = WorldPacketManager(temp_ref=temp_ref, reader=reader, writer=writer)
peername = writer.get_extra_info('peername')
Logger.debug('[World Server]: Accept connection from {}'.format(peername))
Logger.info('[World Server]: trying to process auth session')
auth = AuthManager(reader, writer, temp_ref=temp_ref, world_packet_manager=world_packet_manager)
await auth.process(step=AuthStep.SECOND)
self._register_tasks()
while True:
try:
request = await asyncio.wait_for(reader.read(4096), timeout=1.0)
if request:
response = await asyncio.wait_for(world_packet_manager.process(request), timeout=1.0)
if response:
for packet in response:
writer.write(packet)
await writer.drain()
except TimeoutError:
continue
except Exception as e:
Logger.error('[World Server]: exception, {}'.format(e))
traceback.print_exc()
break
writer.close()
def make(data, *, limit=None):
if limit is not None:
r = StreamReader(loop=event_loop, limit=limit)
else:
r = StreamReader(loop=event_loop)
w = mock_stream_writers()
r.feed_data(data)
r.feed_eof()
return r, w
async def run(self):
self._run = True
started = False
# start with the login procedure
# TODO: this is broken, it fails to restart properly
ctr = 0
while self._run:
if self.protocol and not self.protocol.is_connected:
# reset everything
self.log('we have been disconnected - reset everything')
self.reader = None
self.writer = None
self.protocol = None
self.transport = None
if not self.reader:
self.log('create a StreamReader')
self.reader = StreamReader(loop=self.loop)
if not self.protocol:
self.log('create OWNProtocol')
self.protocol = OWNProtocol(self.reader, log=self.log)
if not self.protocol.is_connected:
if not self._auto_restart:
if started:
self.log("no autorestart... bailing out")
break
started = True
self.log('attempt connection')
self.is_ready.clear()
self.log('resetting the msg handler')
self.msg_handler = self.state_start
try:
self.log('create a new transport connection')
self.transport, _ = await self.loop.create_connection(
lambda: self.protocol, self.host, self.port)
except gaierror:
self.log('sleep about 5 seconds')
await asyncio.sleep(5)
else:
self.log('create a new StreamWriter')
self.writer = StreamWriter(
self.transport, self.protocol, self.reader, self.loop)
self.log('connection is up and running')
try:
pkt = await asyncio.wait_for(self.reader.readuntil(b'##'), timeout=1)
except asyncio.TimeoutError:
pass
else:
# from now on, packets are strings
msg = pkt.decode('ascii')
# self.log('packet (%d)=> %s' % (ctr, msg))
self.log('<= %s' % (msg))
await self.msg_handler(msg)
ctr += 1
self.transport.close()
self.log('AsyncIOOWNConnection.run : %s the end' % (str(self)))
def column_options():
reader = BufferedReader(StreamReader(), constants.BUFFER_SIZE)
writer = BufferedWriter()
context = Context()
context.client_settings = {
"strings_as_bytes": False,
"strings_encoding": constants.STRINGS_ENCODING,
}
column_options = {"reader": reader, "writer": writer, "context": context}
return column_options
async def test_connection_init(event_loop, mock_stream_writers):
r = StreamReader(loop=event_loop)
w = mock_stream_writers()
connection = Connection(r, w)
assert connection.reader is r
assert connection.writer is w
# Did we screw up __str__?
str(connection)
connection.add_header("foo")
str(connection)
def get_connection_mock(ip, port) -> Tuple[MagicMock, MagicMock]:
reader = StreamReader()
reader_mock, writer_mock = (
MagicMock(spec_set=StreamReader, wraps=reader),
MagicMock(spec_set=StreamWriter),
)
writer_mock.get_extra_info.return_value = (ip, port)
reader_mock.__anext__ = AsyncMock(spec_set=reader.__anext__,
wraps=reader.__anext__)
return reader_mock, writer_mock
async def read_json_line(reader: StreamReader) -> Any:
while True:
line = await reader.readline()
if line:
logging.debug(f"read line from socket: {line!r}")
try:
return json.loads(line.decode("utf-8").strip())
except (ValueError, UnicodeDecodeError):
logging.warning("invalid JSON received")
if reader.at_eof():
raise DisconnectedError()
async def handle_connection(self, reader: StreamReader,
writer: StreamWriter):
self._register_tasks()
temp_ref = TempRef()
world_packet_manager = WorldPacketManager(temp_ref=temp_ref,
reader=reader,
writer=writer)
Logger.info('[World Server]: trying to process auth session')
auth = AuthManager(reader,
writer,
temp_ref=temp_ref,
world_packet_manager=world_packet_manager,
session_keys=self.session_keys)
is_authenticated = await auth.process(step=AuthStep.SECOND)
if is_authenticated:
peer_name = writer.get_extra_info('peername')
Logger.success(
'[World Server]: Accept connection from {}'.format(peer_name))
while True:
try:
request = await asyncio.wait_for(reader.read(4096),
timeout=0.01)
if request:
response = await asyncio.wait_for(
world_packet_manager.process(request),
timeout=0.01)
if response:
for packet in response:
writer.write(packet)
await writer.drain()
except TimeoutError:
pass
except BrokenPipeError:
pass
except Exception as e:
Logger.error('[World Server]: exception, {}'.format(e))
traceback.print_exc()
break
finally:
await asyncio.sleep(0.01)
writer.close()
def __init__(self, *args, **kwargs):
"""
Args:
*args:
**kwargs:
"""
sr = StreamReader()
self._stream_reader_orig = sr
self._stream_reader_wr = weakref.ref(sr)
self._stream_writer = None
self.client = node.NeoNode(self)
self._loop = events.get_event_loop()
super().__init__(sr)
def data_received(self, data):
print('Data received {}'.format(data.decode('latin-1')))
if not self._parser:
self._parser = HttpParser()
parser = self._parser
parser.execute(data, len(data))
if parser.is_headers_complete():
print("Headers\n{}".format(parser.get_headers()))
self._transport.pause_reading()
reader = StreamReader()
reader.set_transport(self._transport)
request = HttpRequest(
reader,
method=parser.get_method(),
url=parser.get_url(),
query_string=parser.get_query_string(),
version=parser.get_version(),
headers=parser.get_headers())
self._request = request
ctx = HttpContext(request, HttpResponse(self._transport))
f = asyncio.async(self._app.on_request(ctx))
if parser.is_message_complete():
print("Message complete")
self._parser = None
self._request = None
self._transport.resume_reading()
async def handle_connection(self, reader: StreamReader,
writer: StreamWriter):
peername = writer.get_extra_info('peername')
Logger.info(
'[Login Server]: Accepted connection from {}'.format(peername))
auth = AuthManager(reader, writer)
while not reader.at_eof():
response = await auth.process()
if not response:
break
Logger.warning('[Login Server]: closing...')
writer.close()
async def get_time(self):
NTP_QUERY = bytearray(48)
NTP_QUERY[0] = 0x1B
addr = socket.getaddrinfo(self.host, 123)[0][-1]
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.setblocking(False)
stream = StreamReader(s)
try:
s.sendto(NTP_QUERY, addr)
msg = await stream.read(48)
finally:
await stream.wait_closed()
val = struct.unpack("!I", msg[40:44])[0] - self.ntp_delta
if self.timezone_callback:
val = self.timezone_callback(val)
return val
async def _reader_task(self, reader: StreamReader) -> None:
buffer = b""
while not reader.at_eof():
chunk = await reader.read(1)
if not chunk:
break
buffer += chunk
try:
byte_count, channel_id, frame = unmarshal(buffer)
except UnmarshalingException:
continue
else:
buffer = b""
_logger.debug(f"<- {frame.name} {channel_id}")
await self.dispatch_frame(frame, channel_id)
async def process_request(reader: StreamReader, writer: StreamWriter,
world_packet_mgr: WorldPacketManager):
request = await asyncio.wait_for(
reader.read(4096), timeout=Config.Realm.Settings.min_timeout)
if request:
opcode, data = request[:1], request[1:]
if data:
response = await asyncio.wait_for(
world_packet_mgr.process(opcode=opcode, data=data),
timeout=Config.Realm.Settings.min_timeout)
if response:
for packet in response:
writer.write(packet)
await writer.drain()
async def process_request(reader: StreamReader, writer: StreamWriter, world_packet_mgr: WorldPacketManager):
request: bytes = await wait_for(reader.read(4096), timeout=Config.Realm.Settings.min_timeout)
if request:
size, opcode, data = request[:2], request[2:6], request[6:]
response: List[bytes] = await wait_for(
world_packet_mgr.process(
size=size,
opcode=opcode,
data=data
),
timeout=Config.Realm.Settings.min_timeout
)
if response:
for packet in response:
writer.write(packet)
await writer.drain()
async def handle_connection(self, reader: StreamReader,
writer: StreamWriter):
peername = writer.get_extra_info('peername')
Logger.debug(
'[World Server]: Accept connection from {}'.format(peername))
await self.accept_connection(peername)
auth = AuthSessionManager(reader, writer)
auth.prepare()
await auth.process()
# If I understand it correctly, after SMSG_AUTH_RESPONSE server should waiting for CMSG_CHAR_ENUM
# But next request never sends by client. Client stucks on 'Connected'
try:
next_request = await asyncio.wait_for(reader.read(1024),
timeout=1.0)
Logger.debug(
'[World Server]: next request = {}'.format(next_request))
except TimeoutError:
Logger.error('[World Server]: next request was not received')
finally:
Logger.warning('[World Server]: closing...')
async def ee4_srv(r: StreamReader, w: StreamWriter):
logger.info(f"connection established")
# create new context for client
ctx = SubscriptionContext(r, w)
try:
while not r.exception():
msg_raw = await r.readline()
msg = msg_raw.decode().strip()
logger.info(f"< {msg}")
out = cmd_handler(msg, ctx)
logger.info(f"> {out}")
w.write(f"{out}\r\n".encode())
await w.drain()
except UnicodeDecodeError:
logger.warning("invalid char encountered")
except AttributeError:
logger.warning("attribute error encountered")
except ConnectionError:
logger.warning("connection error encountered")
finally:
logger.info("cleaning up resources")
ctx.end()
logger.info(f"resources purged. connection terminated.")
async def initial_connection_check(self, reader: StreamReader,
writer: StreamWriter):
host_port = writer.get_extra_info('peername')
new_user: Optional[User] = None
try:
# 1. send plain message
writer.write(b'hello')
await writer.drain()
# 2. receive other's header
try:
received = await asyncio.wait_for(reader.read(BUFFER_SIZE),
5.0)
if len(received) == 0:
raise PeerToPeerError('empty msg receive')
header = json.loads(received.decode())
except asyncio.TimeoutError:
raise PeerToPeerError('timeout on other\'s header receive')
except json.JSONDecodeError:
raise PeerToPeerError(
'json decode error on other\'s header receive')
# 3. generate new user
user_header = UserHeader(**header)
new_user = User(user_header, self.number, reader, writer,
host_port, AESCipher.create_key(), SERVER_SIDE)
self.number += 1
if new_user.header.name == V.SERVER_NAME:
raise ConnectionAbortedError('Same origin connection')
# 4. send my public key
my_sec, my_pub = generate_keypair()
send = json.dumps({'public-key': my_pub}).encode()
await new_user.send(send)
self.traffic.put_traffic_up(send)
# 5. receive public key
try:
receive = await new_user.recv()
self.traffic.put_traffic_down(receive)
if len(receive) == 0:
raise ConnectionAbortedError('received msg is zero.')
data = json.loads(receive.decode())
except asyncio.TimeoutError:
raise PeerToPeerError('timeout on public key receive')
except json.JSONDecodeError:
raise PeerToPeerError(
'json decode error on public key receive')
# 6. encrypt and send AES key and header
send = json.dumps({
'aes-key': new_user.aeskey,
'header': self.get_server_header(),
})
key = generate_shared_key(my_sec, data['public-key'])
encrypted = AESCipher.encrypt(key, send.encode())
await new_user.send(encrypted)
self.traffic.put_traffic_up(encrypted)
# 7. receive accept signal
try:
encrypted = await new_user.recv()
self.traffic.put_traffic_down(encrypted)
except asyncio.TimeoutError:
raise PeerToPeerError('timeout on accept signal receive')
receive = AESCipher.decrypt(new_user.aeskey, encrypted)
if receive != b'accept':
raise PeerToPeerError(f"Not accept signal! {receive}")
# 8. accept connection
log.info(
f"established connection as server from {new_user.header.name} {new_user.get_host_port()}"
)
asyncio.ensure_future(self.receive_loop(new_user))
# server port's reachable check
await asyncio.sleep(1.0)
await self.check_reachable(new_user)
return
except (ConnectionAbortedError, ConnectionResetError) as e:
msg = f"disconnect error {host_port} {e}"
except PeerToPeerError as e:
msg = f"peer2peer error {host_port} {e}"
except Exception as e:
msg = "InitialConnCheck: {}".format(e)
log.error(msg, exc_info=True)
# EXCEPTION!
if new_user:
# remove user
self.remove_connection(new_user, msg)
else:
# close socket
log.debug(msg)
try:
writer.write(msg.encode())
await writer.drain()
except Exception:
pass
try:
writer.close()
except Exception:
pass
def __init__(self, data, limit, leave_open=False):
self._reader = StreamReader(limit)
self._reader.feed_data(data)
if not leave_open:
self._reader.feed_eof()
self._mock_write_data = b""
def setUp(self):
self.proxy = Proxy('127.0.0.1', '80', timeout=0.1)
self.proxy._reader['conn'] = StreamReader()
async def test_connection_init(event_loop, mock_stream_writers):
r = StreamReader(loop=event_loop)
w = mock_stream_writers()
connection = Connection(r, w)
assert connection.reader is r
assert connection.writer is w
class HTTPProtocol(asyncio.Protocol):
request_cls = Request
def __init__(self, app=None):
self.app = app
self.router = app.router
self.transport = None
self.reader = None
self.writer = None
self._request: Request = None
self.keep_alive = False
self.loop = asyncio.get_event_loop()
@property
def request(self):
if self._request is None:
self._request = self.request_cls()
self._request.remote_addr = self.remoteaddr
return self._request
@cachedproperty
def remoteaddr(self):
return '{}:{}'.format(*self.transport.get_extra_info('peername'))
def _del_req(self):
self._request = None
def connection_made(self, transport):
self.reader = StreamReader(loop=self.loop)
self.transport = transport
self.request.remote_addr = self.remoteaddr
self.reader.set_transport(transport)
self.loop.create_task(self.parse())
def data_received(self, data):
self.reader.feed_data(data)
# noinspection PyProtectedMember
async def parse(self):
while True:
req = self.request
line = await self.reader.readline()
if not req._f:
req._f = True
_ = _parse_first_line(line)
if not _ or len(_) != 4:
req.parse_error = BadRequestError()
return await self.finish()
req.method, req.path, req.query_string, req.hv = _
else:
if req._body_length:
req._body_lines.append(line)
req._body_length += len(line)
if req._body_length > req.content_length + 1:
raise BadRequestError
elif req._body_length == req.content_length + 1:
return await self.finish()
else:
if line == b'\r\n':
if req.method in {'GET', 'HEAD'
} or not req.content_length:
return await self.finish()
req._body_length = 1
req._body_lines = []
else:
error = _parse_header_line(line, req)
if error is not None:
req.parse_error = error
return await self.finish()
async def finish(self):
req = self.request
delattr(req, '_f')
delattr(req, '_body_length')
self._del_req()
await self.app.handle(req, self)
if req.keepalive:
await self.parse()
else:
self.transport.close()
def send(self, bs):
self.transport.write(bs)
def error(self, err, **kwargs):
handler = self.app.error_handlers.get(err.status)
if handler:
return handler(err, **kwargs)
else:
return err.make_response(**kwargs)
class ControlledConnection:
def __init__(self, limit=100000000):
self._reader = StreamReader(limit)
self._mock_write_data = b""
self._closed = False
self._closed_by_us = False
def _feed_data(self, data):
self._reader.feed_data(data)
def _feed_eof(self):
self._reader.feed_eof()
def _get_mock_write_data(self):
return self._mock_write_data
def _check_eof(self):
if self._reader.at_eof():
self._closed = True
async def read(self, size):
self._check_eof()
try:
return await self._reader.read(min(size, 100))
except Exception:
self._closed = True
raise MalformedDataError
async def readuntil(self, delim):
self._check_eof()
try:
return await self._reader.readuntil(delim)
except Exception:
self._closed = True
raise MalformedDataError
async def readexactly(self, amount):
self._check_eof()
try:
return await self._reader.readexactly(amount)
except Exception:
self._closed = True
raise MalformedDataError
async def write(self, data):
self._mock_write_data += data
return not self._closed
async def add_header(self, header):
pass
def close(self):
self._closed = True
self._closed_by_us = True
async def wait_closed(self):
while not self._closed:
await asyncio.sleep(1)
def closed_by_us(self):
return self._closed_by_us
class NeoProtocol(StreamReaderProtocol):
def __init__(self, *args, quality_check=False, **kwargs):
"""
Args:
*args:
quality_check (bool): there are times when we only establish a connection to check the quality of the node/address
**kwargs:
"""
self._stream_reader = StreamReader()
self._stream_writer = None
nodemanager = kwargs.pop('nodemanager')
self.client = NeoNode(self, nodemanager, quality_check)
self._loop = events.get_event_loop()
super().__init__(self._stream_reader)
def connection_made(self,
transport: asyncio.transports.BaseTransport) -> None:
super().connection_made(transport)
self._stream_writer = StreamWriter(transport, self,
self._stream_reader, self._loop)
if self.client:
asyncio.create_task(self.client.connection_made(transport))
def connection_lost(self, exc: Optional[Exception] = None) -> None:
if self.client:
task = asyncio.create_task(self.client.connection_lost(exc))
task.add_done_callback(
lambda args: super(NeoProtocol, self).connection_lost(exc))
else:
super().connection_lost(exc)
def eof_received(self) -> bool:
self._stream_reader.feed_eof()
self.connection_lost()
return True
# False == Do not keep connection open, this makes sure that `connection_lost` gets called.
# return False
async def send_message(self, message: Message) -> None:
try:
self._stream_writer.write(message.to_array())
await self._stream_writer.drain()
except ConnectionResetError:
# print("connection reset")
self.connection_lost(ConnectionResetError)
except ConnectionError:
# print("connection error")
self.connection_lost(ConnectionError)
except asyncio.CancelledError:
# print("task cancelled, closing connection")
self.connection_lost(asyncio.CancelledError)
except Exception as e:
# print(f"***** woah what happened here?! {traceback.format_exc()}")
self.connection_lost()
async def read_message(self, timeout: int = 30) -> Message:
if timeout == 0:
# avoid memleak. See: https://bugs.python.org/issue37042
timeout = None
async def _read():
try:
message_header = await self._stream_reader.readexactly(24)
magic, command, payload_length, checksum = struct.unpack(
'I 12s I I',
message_header) # uint32, 12byte-string, uint32, uint32
payload_data = await self._stream_reader.readexactly(
payload_length)
payload, = struct.unpack('{}s'.format(payload_length),
payload_data)
except Exception:
# ensures we break out of the main run() loop of Node, which triggers a disconnect callback to clean up
self.client.disconnecting = True
return None
m = Message(magic,
command.rstrip(b'\x00').decode('utf-8'), payload)
if checksum != m.get_checksum(payload):
logger.debug("Message checksum incorrect")
return None
else:
return m
try:
return await asyncio.wait_for(_read(), timeout)
except Exception:
return None
def disconnect(self) -> None:
if self._stream_writer:
self._stream_writer.close()
def stream_reader_writer_slots_before_units(hass):
"""Create a StreamReader and fill it with connection messages."""
reader = StreamReader()
reader.feed_data(
b"\x1b[H\x1b[2JWelcome to the Folding@home Client command server.\n")
reader.feed_data(b"OK\n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"PyON 1 slots\n")
reader.feed_data(
b'[{"id": "00", "status": "READY", "description": "cpu: 3", "options": {"idle": "false"}, "reason": "", "idle": False}]\n'
)
reader.feed_data(b"---\n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"PyON 1 units\n")
reader.feed_data(b"""[{
"id": "00",
"state": "RUNNING",
"error": "NO_ERROR",
"project": 11777,
"run": 0,
"clone": 17250,
"gen": 0,
"core": "0x22",
"unit": "0x00000003287234c95e774a8488c5d4ea",
"percentdone": "72.51%",
"eta": "1 hours 04 mins",
"ppd": "359072",
"creditestimate": "58183",
"waitingon": "",
"nextattempt": "0.00 secs",
"timeremaining": "8.08 days",
"totalframes": 100,
"framesdone": 72,
"assigned": "2020-03-28T08: 33: 55Z",
"timeout": "2020-03-29T08: 33: 55Z",
"deadline": "2020-04-05T13: 21: 54Z",
"ws": "40.114.52.201",
"cs": "13.82.98.119",
"attempts": 0,
"slot": "00",
"tpf": "2 mins 20 secs",
"basecredit": "9405"
}
]\n""")
reader.feed_data(b"---\n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_data(b"> \n")
reader.feed_eof()
stream_writer = AsyncMock()
stream_writer.close = MagicMock()
return (reader, stream_writer)
def __init__(self, limit=100000000):
self._reader = StreamReader(limit)
self._mock_write_data = b""
self._closed = False
self._closed_by_us = False
def proxy():
proxy = Proxy('127.0.0.1', '80', timeout=0.1)
proxy._reader['conn'] = StreamReader()
return proxy
class AsyncIOOWNConnection(object):
SOCKET_MODES = ('*99*0##', '*99*1##', '*99*9##')
MODE_COLORS = (COLOR_LT_BLUE, COLOR_LT_GREEN, COLOR_LT_CYAN)
MODE_NAMES = ('CMD', 'MON', 'SCMD')
ACK = '*#*1##'
def __init__(self, host, port, passwd, queue, mode,
log=None, loop=None):
self.host = host
self.port = port
self.passwd = passwd
self.queue = queue
self.mode = mode
self.loop = loop
self._auto_restart = True
self._run = False
if self.loop is None:
self.loop = asyncio.get_event_loop()
self.is_ready = asyncio.Event(loop=self.loop)
if log is None:
hdr = '[%s:%d %s]' % \
(self.host, self.port, self.MODE_NAMES[self.mode])
self.log = \
get_logger(header=hdr,
color=self.MODE_COLORS[self.mode])
else:
self.log = log
self.reader = None
self.writer = None
self.protocol = None
self.transport = None
@property
def auto_restart(self):
return self._auto_restart
@auto_restart.setter
def auto_restart(self, value):
self._auto_restart = value
async def run(self):
self._run = True
started = False
# start with the login procedure
# TODO: this is broken, it fails to restart properly
ctr = 0
while self._run:
if self.protocol and not self.protocol.is_connected:
# reset everything
self.log('we have been disconnected - reset everything')
self.reader = None
self.writer = None
self.protocol = None
self.transport = None
if not self.reader:
self.log('create a StreamReader')
self.reader = StreamReader(loop=self.loop)
if not self.protocol:
self.log('create OWNProtocol')
self.protocol = OWNProtocol(self.reader, log=self.log)
if not self.protocol.is_connected:
if not self._auto_restart:
if started:
self.log("no autorestart... bailing out")
break
started = True
self.log('attempt connection')
self.is_ready.clear()
self.log('resetting the msg handler')
self.msg_handler = self.state_start
try:
self.log('create a new transport connection')
self.transport, _ = await self.loop.create_connection(
lambda: self.protocol, self.host, self.port)
except gaierror:
self.log('sleep about 5 seconds')
await asyncio.sleep(5)
else:
self.log('create a new StreamWriter')
self.writer = StreamWriter(
self.transport, self.protocol, self.reader, self.loop)
self.log('connection is up and running')
try:
pkt = await asyncio.wait_for(self.reader.readuntil(b'##'), timeout=1)
except asyncio.TimeoutError:
pass
else:
# from now on, packets are strings
msg = pkt.decode('ascii')
# self.log('packet (%d)=> %s' % (ctr, msg))
self.log('<= %s' % (msg))
await self.msg_handler(msg)
ctr += 1
self.transport.close()
self.log('AsyncIOOWNConnection.run : %s the end' % (str(self)))
def stop(self):
self.log('stop requested', LOG_DEBUG)
self._run = False
async def send_packet(self, msg):
self.log('=> %s' % (msg))
if isinstance(msg, str):
pkt = msg.encode('ascii')
self.writer.write(pkt)
await self.writer.drain()
async def state_start(self, msg):
if msg == self.ACK:
self.msg_handler = self.state_login
cmd_msg = self.SOCKET_MODES[self.mode]
await self.send_packet(cmd_msg)
else:
self.log('we didn\'t get ACK')
async def state_login(self, msg):
if msg == '*98*2##':
# this is a call for sha2 hmac authentication
self.msg_handler = self.state_hmac_sha2
await self.send_packet(self.ACK)
else:
# attempt matching the old password system
ops_m = re.match(r'^\*#(\d+)##$', msg)
if ops_m is not None:
nonce = ops_m.groups()[0]
# calculate the password
passwd = ownCalcPass(self.passwd, nonce)
passwd_msg = ('*#%s##' % (passwd))
self.msg_handler = self.state_auth
await self.send_packet(passwd_msg)
else:
self.log("unable to parse the openpassword nonce request")
await asyncio.sleep(0)
async def state_hmac_sha2(self, msg):
ra_m = re.match(r'^\*#(\d{128})##$', msg)
if ra_m is not None:
ra = ra_m.group(1)
hmac = ownCalcHmacSha2(self.passwd, ra)
rb, hmac_client, self.hmac_server = hmac
self.msg_handler = self.state_hmac_sha2_check_response
hmac_packet = ('*#%s*%s##' % (rb, hmac_client))
await self.send_packet(hmac_packet)
else:
self.log("unable to parse the hmac_sha2 request")
await asyncio.sleep(0)
async def state_hmac_sha2_check_response(self, msg):
c_resp_m = re.match(r'^\*#(\d{128})##$', msg)
if c_resp_m is not None:
c_resp = c_resp_m.group(1)
if c_resp == self.hmac_server:
self.is_ready.set()
self.msg_handler = self.state_dispatch
await self.send_packet(self.ACK)
else:
self.log("wrong response from server, expected %s", self.hmac_server)
await asyncio.sleep(0)
else:
self.log("unable to parse the hmac_sha2 request")
await asyncio.sleep(0)
async def state_auth(self, msg):
if msg == self.ACK:
self.is_ready.set()
self.msg_handler = self.state_dispatch
async def state_dispatch(self, msg):
await self.queue.put((msg, self.mode, ))
