burmatovan/openuds
1
0
Fork 0
forked from shaba/openuds
Code Issues Pull Requests Projects Releases Wiki Activity

Added end notification to UDS

Browse Source
This commit is contained in:
Adolfo Gómez García 2022-02-01 13:49:34 +01:00
parent 4893fa4fbe
commit 9041e48080
1 changed files with 16 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
tunnel-server/src/uds_tunnel/tunnel.py
View File

@ -26,6 +26,8 @@ class TunnelProtocol(asyncio.Protocol):
runner: typing.Any runner: typing.Any
# Command buffer # Command buffer
cmd: bytes cmd: bytes
# Ticket
notify_ticket: bytes
# owner Proxy class # owner Proxy class
owner: 'proxy.Proxy' owner: 'proxy.Proxy'
# source of connection # source of connection
@ -55,6 +57,7 @@ class TunnelProtocol(asyncio.Protocol):
# transport is undefined until connection_made is called # transport is undefined until connection_made is called
self.finished = asyncio.Future() self.finished = asyncio.Future()
self.cmd = b'' self.cmd = b''
self.notify_ticket = b''
self.owner = owner self.owner = owner
self.source = ('', 0) self.source = ('', 0)
self.destination = ('', 0) self.destination = ('', 0)
@ -85,6 +88,7 @@ class TunnelProtocol(asyncio.Protocol):
# store for future use # store for future use
self.destination = (result['host'], int(result['port'])) self.destination = (result['host'], int(result['port']))
self.notify_ticket = result['notify'].encode()
logger.info( logger.info(
'OPEN TUNNEL FROM %s to %s', 'OPEN TUNNEL FROM %s to %s',
@ -193,6 +197,15 @@ class TunnelProtocol(asyncio.Protocol):
logger.debug('Data received: %s', len(data)) logger.debug('Data received: %s', len(data))
self.runner(data) # send data to current runner (command or proxy) self.runner(data) # send data to current runner (command or proxy)
def notifyEnd(self):
if self.notify_ticket:
asyncio.get_event_loop().create_task(
TunnelProtocol.notifyEndToUds(
self.owner.cfg, self.notify_ticket, self.stats_manager
)
)
self.notify_ticket = b'' # Clean up so no more notifications
def connection_lost(self, exc: typing.Optional[Exception]) -> None: def connection_lost(self, exc: typing.Optional[Exception]) -> None:
logger.debug('Connection closed : %s', exc) logger.debug('Connection closed : %s', exc)
self.finished.set_result(True) self.finished.set_result(True)
@ -200,6 +213,7 @@ class TunnelProtocol(asyncio.Protocol):
self.other_side.transport.close() self.other_side.transport.close()
else: else:
self.stats_manager.close() self.stats_manager.close()
self.notifyEnd()
# helpers # helpers
# source address, pretty format # source address, pretty format
@ -221,6 +235,8 @@ class TunnelProtocol(asyncio.Protocol):
self.stats_manager.recv, self.stats_manager.recv,
int(self.stats_manager.end - self.stats_manager.start), int(self.stats_manager.end - self.stats_manager.start),
) )
# Notify end to uds
self.notifyEnd()
else: else:
logger.info('TERMINATED %s', self.pretty_source()) logger.info('TERMINATED %s', self.pretty_source())