diff --git a/core/controllers/profile_state/data_filters.py b/core/controllers/profile_state/data_filters.py index 8ab7617..c5d01fc 100644 --- a/core/controllers/profile_state/data_filters.py +++ b/core/controllers/profile_state/data_filters.py @@ -27,16 +27,10 @@ Inputs: def get_supported_protocols_for_location(location_codes: str, profile_type: str = None) -> Result: country, city = location_codes.split("_") - print(f"country is {country}") - print(f"city is {city}") # this IS doing sql: location = get_profile_location_data(country_code=country, city_code=city) - print(f"location is {location}") - print(f"location type is {type(location)}") - # return location - if not location: error_msg = f"This is a developer error. No location data found for {location_codes}" logger.error(error_msg) @@ -58,10 +52,7 @@ def _get_supported_protocols_for_location(location: Location, as_strings: bool = for each_protocol in ALL_PROTOCOLS: supported = location.supports_protocol(enum_value=each_protocol) - print(f"supported for {each_protocol} is {supported}") - print(f"type of supported is {type(supported)}") if supported: - print(f"THIS IS TRUE for {each_protocol}") SUPPORTED_PROTOCOLS.append(each_protocol) return SUPPORTED_PROTOCOLS diff --git a/core/controllers/profile_state/update_profile.py b/core/controllers/profile_state/update_profile.py index a4c1997..0cbf204 100644 --- a/core/controllers/profile_state/update_profile.py +++ b/core/controllers/profile_state/update_profile.py @@ -14,7 +14,7 @@ from core.models.BaseProfile import get_application_version # from core.models.Subscription import Subscription # from core.models.session.SessionProfile import SessionProfile # from core.models.system.SystemProfile import SystemProfile -# from core.models.orm_models.Location import Location +from core.models.orm_models.Location import Location # from core.models.orm_models.Operator import Operator SYSTEMWIDE_CHOICES = ['wireguard', 'hysteria2', 'vless'] @@ -93,13 +93,16 @@ def update_profile( # ============= LOCATION ============= elif key == "location": country_code, city_code = new_value.split("_") - location = get_profile_location_data( # SQLAlchemy + new_location = get_profile_location_data( # SQLAlchemy country_code=country_code, city_code=city_code ) - # SQLAlchemy foreign key assignment — fills in id, timezone, operator, etc. - profile.location = location - + pre_existing_protocol = profile.connection.code + if new_location.supports_protocol(pre_existing_protocol): + profile.location = new_location + else: + return Result(valid=False, message=f"The {new_location.name} location doesn't currently support the {str(pre_existing_protocol)} protocol", error_type=ResultError.WRONG_LOCATION) + # Subscription gets wiped on a location change, and is outside ORM profile.subscription = None else: @@ -117,12 +120,23 @@ def update_profile( def change_protocol(profile: Profile, new_value: str) -> Result: logger.info(f"changing the protcol to {new_value}") + # ============= SAME VALUE ============= if profile.connection.code == new_value: error_msg = f"This is NOT changing the protocol, it already was {new_value}" logger.error(error_msg) return Result(valid=False, message=error_msg, error_type=ResultError.INVALID_INPUT) + # ============= LOCATION SUPPORTS PROTOCOL ============= + logger.info("checking protocol support for this location..") + pre_existing_location = profile.location + + if not pre_existing_location.supports_protocol(new_value): + logger.error(f"User is trying to update the protocol to {new_value}, but selected an invalid location to do so.") + error_msg = f"This location does not yet support the {str(new_value)} protocol." + return Result(valid=False, message=error_msg, error_type=ResultError.WRONG_LOCATION) + logger.info("Passed location check") + # ============= SYSTEMWIDE ============= logger.info(f"Evaluating the profile's connection type of {profile.type}") if profile.type == ProfileType.SYSTEM: diff --git a/core/models/orm_models/Location.py b/core/models/orm_models/Location.py index c2f7b34..26d449a 100644 --- a/core/models/orm_models/Location.py +++ b/core/models/orm_models/Location.py @@ -56,9 +56,9 @@ class Location(BaseModel): - def supports_protocol(self, enum_value: ProtocolChoice) -> bool: - # print(f"inside {self.country_name} {self.name} location model's supports_protocol, the submitted protocol is {enum_value}") - print(f"enum_value type is {type(enum_value)}") + def supports_protocol(self, enum_value: ProtocolChoice|str) -> bool: + """Check if a given location supports a given protocol""" + if enum_value is None: return False @@ -72,7 +72,6 @@ class Location(BaseModel): string_version = enum_value field_name = f'is_{string_version}_capable' - print(f"field_name is {field_name}") return bool(getattr(self, field_name))