Improve exception override-related logic

This commit is contained in:
codeking 2025-11-04 20:48:01 +01:00
parent b1877e47a2
commit 5080014070

View file

@ -1,5 +1,5 @@
from core.Constants import Constants
from core.Errors import MissingSubscriptionError, InvalidSubscriptionError, UnknownConnectionTypeError
from core.Errors import MissingSubscriptionError, InvalidSubscriptionError, UnknownConnectionTypeError, ConnectionUnprotectedError, ProfileStateConflictError
from core.controllers.ApplicationController import ApplicationController
from core.controllers.ApplicationVersionController import ApplicationVersionController
from core.controllers.ClientController import ClientController
@ -95,8 +95,9 @@ if __name__ == '__main__':
pristine_parser = argparse.ArgumentParser(add_help=False)
pristine_parser.add_argument('--pristine', '-p', action='store_true')
safe_parser = argparse.ArgumentParser(add_help=False)
safe_parser.add_argument('--force', '-f', action='store_true')
security_parser = argparse.ArgumentParser(add_help=False)
security_parser.add_argument('--disable-connection-protection', action='store_true')
security_parser.add_argument('--disable-profile-state-protection', action='store_true')
main_parser = argparse.ArgumentParser(prog=__get_name())
main_parser.add_argument('--version', '-v', action='version', version=f'{__get_name()} v{__get_version()}')
@ -115,7 +116,7 @@ if __name__ == '__main__':
profile_create_parser = profile_subparsers.add_parser('create')
profile_create_subparsers = profile_create_parser.add_subparsers(title='profile_types', dest='profile_type')
session_profile_create_parser = profile_create_subparsers.add_parser('session', parents=[profile_base_parser, safe_parser])
session_profile_create_parser = profile_create_subparsers.add_parser('session', parents=[profile_base_parser])
session_profile_create_parser.add_argument('--name', '-n', default='')
session_profile_create_parser.add_argument('--location', '-l', default='')
@ -124,16 +125,16 @@ if __name__ == '__main__':
session_profile_create_parser.add_argument('--mask-connection', '-m', action='store_true')
session_profile_create_parser.add_argument('--resolution', '-r', default='1280x720')
system_profile_create_parser = profile_create_subparsers.add_parser('system', parents=[profile_base_parser, safe_parser])
system_profile_create_parser = profile_create_subparsers.add_parser('system', parents=[profile_base_parser])
system_profile_create_parser.add_argument('--name', '-n', default='')
system_profile_create_parser.add_argument('--location', '-l', default='')
system_profile_create_parser.add_argument('--connection', '-c', dest='connection_type', choices=['wireguard'], default='wireguard')
profile_subparsers.add_parser('destroy', parents=[profile_base_parser, safe_parser])
profile_subparsers.add_parser('destroy', parents=[profile_base_parser])
profile_subparsers.add_parser('enable', parents=[profile_base_parser, pristine_parser, safe_parser])
profile_subparsers.add_parser('disable', parents=[profile_base_parser, safe_parser])
profile_subparsers.add_parser('enable', parents=[profile_base_parser, pristine_parser, security_parser])
profile_subparsers.add_parser('disable', parents=[profile_base_parser, security_parser])
application_parser = main_subparsers.add_parser('application')
application_subparsers = application_parser.add_subparsers(title='subcommands', dest='subcommand')
@ -168,6 +169,16 @@ if __name__ == '__main__':
arguments = main_parser.parse_args()
ignore = []
if getattr(arguments, 'disable_connection_protection', False):
ignore.append(ConnectionUnprotectedError)
if getattr(arguments, 'disable_profile_state_protection', False):
ignore.append(ProfileStateConflictError)
ignore = tuple(ignore)
if arguments.command is None:
main_parser.print_help()
@ -231,7 +242,7 @@ if __name__ == '__main__':
if profile is not None:
try:
ProfileController.enable(profile, force=arguments.force, pristine=arguments.pristine, asynchronous=True, profile_observer=profile_observer, application_version_observer=application_version_observer, connection_observer=connection_observer)
ProfileController.enable(profile, ignore=ignore, pristine=arguments.pristine, asynchronous=True, profile_observer=profile_observer, application_version_observer=application_version_observer, connection_observer=connection_observer)
except (InvalidSubscriptionError, MissingSubscriptionError) as exception:
@ -279,7 +290,7 @@ if __name__ == '__main__':
else:
raise RuntimeError('The subscription could not be activated. Please try again later.')
ProfileController.enable(profile, force=arguments.force, pristine=arguments.pristine, asynchronous=True, profile_observer=profile_observer, application_version_observer=application_version_observer, connection_observer=connection_observer)
ProfileController.enable(profile, ignore=ignore, pristine=arguments.pristine, asynchronous=True, profile_observer=profile_observer, application_version_observer=application_version_observer, connection_observer=connection_observer)
elif manage_subscription_input == '2':
@ -291,7 +302,7 @@ if __name__ == '__main__':
if subscription is not None:
ProfileController.attach_subscription(profile, subscription)
ProfileController.enable(profile, force=arguments.force, pristine=arguments.pristine, asynchronous=True, profile_observer=profile_observer, application_version_observer=application_version_observer, connection_observer=connection_observer)
ProfileController.enable(profile, ignore=ignore, pristine=arguments.pristine, asynchronous=True, profile_observer=profile_observer, application_version_observer=application_version_observer, connection_observer=connection_observer)
else:
@ -312,7 +323,7 @@ if __name__ == '__main__':
profile = ProfileController.get(arguments.id)
if profile is not None:
ProfileController.disable(profile, force=arguments.force, profile_observer=profile_observer)
ProfileController.disable(profile, ignore=ignore, profile_observer=profile_observer)
else:
main_parser.error('the following argument should be a valid reference: --id/-i')