diff --git a/server/src/uds/core/transports/transport.py b/server/src/uds/core/transports/transport.py index 6e12aade..ffdca8c0 100644 --- a/server/src/uds/core/transports/transport.py +++ b/server/src/uds/core/transports/transport.py @@ -55,9 +55,12 @@ logger = logging.getLogger(__name__) DIRECT_GROUP = _('Direct') TUNNELED_GROUP = _('Tunneled') + class TransportScript(typing.NamedTuple): script: str = '' - script_type: typing.Union[typing.Literal['python'], typing.Literal['lua']] = 'python' # currently only python is supported + script_type: typing.Union[ + typing.Literal['python'], typing.Literal['lua'] + ] = 'python' # currently only python is supported signature_b64: str = '' # Signature of the script in base64 parameters: typing.Mapping[str, typing.Any] = {} @@ -66,7 +69,10 @@ class TransportScript(typing.NamedTuple): """ Returns encoded parameters for transport script """ - return codecs.encode(codecs.encode(json.dumps(self.parameters).encode(), 'bz2'), 'base64').decode() + return codecs.encode( + codecs.encode(json.dumps(self.parameters).encode(), 'bz2'), 'base64' + ).decode() + class Transport(Module): """ @@ -270,7 +276,9 @@ class Transport(Module): logger.debug('Transport script: %s', transport_script) return TransportScript( - script=codecs.encode(codecs.encode(transport_script.script.encode(), 'bz2'), 'base64') + script=codecs.encode( + codecs.encode(transport_script.script.encode(), 'bz2'), 'base64' + ) .decode() .replace('\n', ''), signature_b64=transport_script.signature_b64, @@ -279,7 +287,7 @@ class Transport(Module): def getRelativeScript( self, scriptName: str, params: typing.Mapping[str, typing.Any] - ) -> typing.Tuple[str, str, typing.Mapping[str, typing.Any]]: + ) -> 'TransportScript': """Returns a script that will be executed on client, but will be downloaded from server Args: @@ -287,16 +295,26 @@ class Transport(Module): params: Parameters for the return tuple """ # Reads script and signature - basePath = os.path.dirname(sys.modules[self.__module__].__file__ or 'not_found') # Will raise if not found + basePath = os.path.dirname( + sys.modules[self.__module__].__file__ or 'not_found' + ) # Will raise if not found script = open(os.path.join(basePath, scriptName), 'r').read() signature = open(os.path.join(basePath, scriptName + '.signature'), 'r').read() - return script, signature, params + return TransportScript( + script=script, + script_type='python', + signature_b64=signature, + parameters=params, + ) def getScript( - self, osName: str, type: typing.Literal['tunnel', 'direct'], params: typing.Mapping[str, typing.Any] - ) -> typing.Tuple[str, str, typing.Mapping[str, typing.Any]]: + self, + osName: str, + type: typing.Literal['tunnel', 'direct'], + params: typing.Mapping[str, typing.Any], + ) -> 'TransportScript': """ Returns a script for the given os and type """ diff --git a/server/src/uds/transports/RDP/rdp_base.py b/server/src/uds/transports/RDP/rdp_base.py index e993d19c..ac5e1802 100644 --- a/server/src/uds/transports/RDP/rdp_base.py +++ b/server/src/uds/transports/RDP/rdp_base.py @@ -433,25 +433,3 @@ class BaseRDPTransport(transports.Transport): password, altUsername=username, ) - - def getScript( - self, - scriptNameTemplate: str, - osName: str, - params: typing.Mapping[str, typing.Any], - ) -> transports.TransportScript: - # Reads script - scriptNameTemplate = scriptNameTemplate.format(osName) - with open(os.path.join(os.path.dirname(__file__), scriptNameTemplate)) as f: - script = f.read() - # Reads signature - with open( - os.path.join(os.path.dirname(__file__), scriptNameTemplate + '.signature') - ) as f: - signature = f.read() - return transports.TransportScript( - script=script, - script_type='python', - signature_b64=signature, - parameters=params, - ) diff --git a/server/src/uds/transports/SPICE/spice_base.py b/server/src/uds/transports/SPICE/spice_base.py index 40a07a72..22813009 100644 --- a/server/src/uds/transports/SPICE/spice_base.py +++ b/server/src/uds/transports/SPICE/spice_base.py @@ -201,22 +201,3 @@ class BaseSpiceTransport(transports.Transport): password: str, ) -> typing.Dict[str, str]: return self.processUserPassword(userService, user, password) - - def getScript( - self, scriptNameTemplate: str, osName: str, params: typing.Dict[str, typing.Any] - ) -> transports.TransportScript: - # Reads script - scriptNameTemplate = scriptNameTemplate.format(osName) - with open(os.path.join(os.path.dirname(__file__), scriptNameTemplate)) as f: - script = f.read() - # Reads signature - with open( - os.path.join(os.path.dirname(__file__), scriptNameTemplate + '.signature') - ) as f: - signature = f.read() - return transports.TransportScript( - script=script, - script_type='python', - signature_b64=signature, - parameters=params, - ) diff --git a/server/src/uds/transports/X2GO/x2go_base.py b/server/src/uds/transports/X2GO/x2go_base.py index ce6c33b5..9d1cff8d 100644 --- a/server/src/uds/transports/X2GO/x2go_base.py +++ b/server/src/uds/transports/X2GO/x2go_base.py @@ -285,22 +285,3 @@ class BaseX2GOTransport(transports.Transport): authScript = self.getAuthorizeScript(userName, pub) userServiceManager().sendScript(userService, authScript) return priv, pub - - def getScript( - self, scriptNameTemplate: str, osName: str, params: typing.Dict[str, typing.Any] - ) -> 'transports.TransportScript': - # Reads script - scriptNameTemplate = scriptNameTemplate.format(osName) - with open(os.path.join(os.path.dirname(__file__), scriptNameTemplate)) as f: - script = f.read() - # Reads signature - with open( - os.path.join(os.path.dirname(__file__), scriptNameTemplate + '.signature') - ) as f: - signature = f.read() - return transports.TransportScript( - script=script, - script_type='python', - signature_b64=signature, - parameters=params, - )