diff --git a/src/falconpy/_auth_object/_uber_interface.py b/src/falconpy/_auth_object/_uber_interface.py index 7ffcbd520..b23e61b11 100644 --- a/src/falconpy/_auth_object/_uber_interface.py +++ b/src/falconpy/_auth_object/_uber_interface.py @@ -150,6 +150,15 @@ def logout(self) -> bool: return bool(result["status_code"] == 200) + def __enter__(self): + """Allow for entry as a context manager.""" + return self + + def __exit__(self, *args): + """Discard our token when we exit the context.""" + self.logout() + return args + # Legacy property getters maintained for backwards functionality. def authenticated(self) -> bool: """Return the current authentication status.""" diff --git a/src/falconpy/_service_class/_service_class.py b/src/falconpy/_service_class/_service_class.py index a11a4df5b..6d98def7c 100644 --- a/src/falconpy/_service_class/_service_class.py +++ b/src/falconpy/_service_class/_service_class.py @@ -230,6 +230,15 @@ def override(self, exp=expand_result )) + def __enter__(self): + """Allow for entry as a context manager.""" + return self + + def __exit__(self, *args): + """Discard our token when we exit the context.""" + self.logout() + return args + # ___ ____ ____ ___ ____ ____ ___ _ ____ ____ # |__] |__/ | | |__] |___ |__/ | | |___ [__ # | | \ |__| | |___ | \ | | |___ ___]