Impersonate start/end (#92)

Maria Korlotian · web-flow · commit 6f97c000e541 · 2020-11-26T15:18:38.000+01:00
* Add start and end impersonation commands

* Unify the method name - rename build to call
diff --git a/castle/api/approve_device.py b/castle/api/approve_device.py
@@ -5,4 +5,4 @@
 class APIApproveDevice(object):
     @staticmethod
     def call(device_token):
-        return APIRequest().call(CommandsApproveDevice.build(device_token))
+        return APIRequest().call(CommandsApproveDevice.call(device_token))
diff --git a/castle/api/get_device.py b/castle/api/get_device.py
@@ -5,4 +5,4 @@
 class APIGetDevice(object):
     @staticmethod
     def call(device_token):
-        return APIRequest().call(CommandsGetDevice.build(device_token))
+        return APIRequest().call(CommandsGetDevice.call(device_token))
diff --git a/castle/api/get_devices_for_user.py b/castle/api/get_devices_for_user.py
@@ -5,4 +5,4 @@
 class APIGetDevicesForUser(object):
     @staticmethod
     def call(user_id):
-        return APIRequest().call(CommandsGetDevicesForUser.build(user_id))
+        return APIRequest().call(CommandsGetDevicesForUser.call(user_id))
diff --git a/castle/api/report_device.py b/castle/api/report_device.py
@@ -5,4 +5,4 @@
 class APIReportDevice(object):
     @staticmethod
     def call(device_token):
-        return APIRequest().call(CommandsReportDevice.build(device_token))
+        return APIRequest().call(CommandsReportDevice.call(device_token))
diff --git a/castle/api/review.py b/castle/api/review.py
@@ -5,4 +5,4 @@
 class APIReview(object):
     @staticmethod
     def call(review_id):
-        return APIRequest().call(CommandsReview.build(review_id))
+        return APIRequest().call(CommandsReview.call(review_id))
diff --git a/castle/client.py b/castle/client.py
@@ -1,7 +1,8 @@
 from castle.api_request import APIRequest
 from castle.commands.authenticate import CommandsAuthenticate
 from castle.commands.identify import CommandsIdentify
-from castle.commands.impersonate import CommandsImpersonate
+from castle.commands.start_impersonation import CommandsStartImpersonation
+from castle.commands.end_impersonation import CommandsEndImpersonation
 from castle.commands.track import CommandsTrack
 from castle.configuration import configuration
 from castle.context.prepare import ContextPrepare
@@ -43,7 +44,7 @@ def _add_timestamp_if_necessary(self, options):
     def authenticate(self, options):
         if self.tracked():
             self._add_timestamp_if_necessary(options)
-            command = CommandsAuthenticate(self.context).build(options)
+            command = CommandsAuthenticate(self.context).call(options)
             try:
                 response = self.api.call(command)
                 response.update(failover=False, failover_reason=None)
@@ -61,11 +62,18 @@ def identify(self, options):
         if not self.tracked():
             return None
         self._add_timestamp_if_necessary(options)
-        return self.api.call(CommandsIdentify(self.context).build(options))
+        return self.api.call(CommandsIdentify(self.context).call(options))
 
-    def impersonate(self, options):
+    def start_impersonation(self, options):
         self._add_timestamp_if_necessary(options)
-        response = self.api.call(CommandsImpersonate(self.context).build(options))
+        response = self.api.call(CommandsStartImpersonation(self.context).call(options))
+        if not response.get('success'):
+            raise ImpersonationFailed
+        return response
+
+    def end_impersonation(self, options):
+        self._add_timestamp_if_necessary(options)
+        response = self.api.call(CommandsEndImpersonation(self.context).call(options))
         if not response.get('success'):
             raise ImpersonationFailed
         return response
@@ -74,7 +82,7 @@ def track(self, options):
         if not self.tracked():
             return None
         self._add_timestamp_if_necessary(options)
-        return self.api.call(CommandsTrack(self.context).build(options))
+        return self.api.call(CommandsTrack(self.context).call(options))
 
     def disable_tracking(self):
         self.do_not_track = True
diff --git a/castle/commands/approve_device.py b/castle/commands/approve_device.py
@@ -5,7 +5,7 @@
 class CommandsApproveDevice(object):
 
     @staticmethod
-    def build(device_token):
+    def call(device_token):
         ValidatorsPresent.call({'device_token': device_token}, 'device_token')
 
         return Command(
diff --git a/castle/commands/authenticate.py b/castle/commands/authenticate.py
@@ -9,7 +9,7 @@ class CommandsAuthenticate(object):
     def __init__(self, context):
         self.context = context
 
-    def build(self, options):
+    def call(self, options):
         ValidatorsPresent.call(options, 'event')
         context = ContextMerge.call(self.context, options.get('context'))
         context = ContextSanitize.call(context)
diff --git a/castle/commands/end_impersonation.py b/castle/commands/end_impersonation.py
@@ -0,0 +1,23 @@
+from castle.command import Command
+from castle.utils.timestamp import UtilsTimestamp as generate_timestamp
+from castle.context.merge import ContextMerge
+from castle.context.sanitize import ContextSanitize
+from castle.validators.present import ValidatorsPresent
+
+
+class CommandsEndImpersonation(object):
+    def __init__(self, context):
+        self.context = context
+
+    def call(self, options):
+        ValidatorsPresent.call(options, 'user_id')
+
+        context = ContextMerge.call(self.context, options.get('context'))
+        context = ContextSanitize.call(context)
+        ValidatorsPresent.call(context, 'user_agent', 'ip')
+
+        if context:
+            options.update({'context': context})
+        options.update({'sent_at': generate_timestamp.call()})
+
+        return Command(method='delete', path='impersonate', data=options)
diff --git a/castle/commands/get_device.py b/castle/commands/get_device.py
@@ -5,7 +5,7 @@
 class CommandsGetDevice(object):
 
     @staticmethod
-    def build(device_token):
+    def call(device_token):
         ValidatorsPresent.call({'device_token': device_token}, 'device_token')
 
         return Command(
diff --git a/castle/commands/get_devices_for_user.py b/castle/commands/get_devices_for_user.py
@@ -5,7 +5,7 @@
 class CommandsGetDevicesForUser(object):
 
     @staticmethod
-    def build(user_id):
+    def call(user_id):
         ValidatorsPresent.call({'user_id': user_id}, 'user_id')
 
         return Command(
diff --git a/castle/commands/identify.py b/castle/commands/identify.py
@@ -9,7 +9,7 @@ class CommandsIdentify(object):
     def __init__(self, context):
         self.context = context
 
-    def build(self, options):
+    def call(self, options):
         ValidatorsNotSupported.call(options, 'properties')
         context = ContextMerge.call(self.context, options.get('context'))
         context = ContextSanitize.call(context)
diff --git a/castle/commands/report_device.py b/castle/commands/report_device.py
@@ -5,7 +5,7 @@
 class CommandsReportDevice(object):
 
     @staticmethod
-    def build(device_token):
+    def call(device_token):
         ValidatorsPresent.call({'device_token': device_token}, 'device_token')
 
         return Command(
diff --git a/castle/commands/review.py b/castle/commands/review.py
@@ -5,7 +5,7 @@
 class CommandsReview(object):
 
     @staticmethod
-    def build(review_id):
+    def call(review_id):
         ValidatorsPresent.call({'review_id': review_id}, 'review_id')
 
         return Command(
diff --git a/castle/commands/start_impersonation.py b/castle/commands/start_impersonation.py
@@ -5,11 +5,11 @@
 from castle.validators.present import ValidatorsPresent
 
 
-class CommandsImpersonate(object):
+class CommandsStartImpersonation(object):
     def __init__(self, context):
         self.context = context
 
-    def build(self, options):
+    def call(self, options):
         ValidatorsPresent.call(options, 'user_id')
 
         context = ContextMerge.call(self.context, options.get('context'))
@@ -20,6 +20,4 @@ def build(self, options):
             options.update({'context': context})
         options.update({'sent_at': generate_timestamp.call()})
 
-        method = ('delete' if options.get('reset', False) else 'post')
-
-        return Command(method=method, path='impersonate', data=options)
+        return Command(method='post', path='impersonate', data=options)
diff --git a/castle/commands/track.py b/castle/commands/track.py
@@ -9,7 +9,7 @@ class CommandsTrack(object):
     def __init__(self, context):
         self.context = context
 
-    def build(self, options):
+    def call(self, options):
         ValidatorsPresent.call(options, 'event')
         context = ContextMerge.call(self.context, options.get('context'))
         context = ContextSanitize.call(context)
diff --git a/castle/test/__init__.py b/castle/test/__init__.py
@@ -16,12 +16,13 @@
     'castle.test.command_test',
     'castle.test.commands.approve_device_test',
     'castle.test.commands.authenticate_test',
+    'castle.test.commands.end_impersonation_test',
     'castle.test.commands.get_device_test',
     'castle.test.commands.get_devices_for_user_test',
     'castle.test.commands.identify_test',
-    'castle.test.commands.impersonate_test',
     'castle.test.commands.report_device_test',
     'castle.test.commands.review_test',
+    'castle.test.commands.start_impersonation_test',
     'castle.test.commands.track_test',
     'castle.test.configuration_test',
     'castle.test.context.get_default_test',
diff --git a/castle/test/client_test.py b/castle/test/client_test.py
@@ -46,7 +46,7 @@ def test_init(self):
         self.assertIsInstance(client.api, APIRequest)
 
     @responses.activate
-    def test_impersonate(self):
+    def test_start_impersonation(self):
         response_text = {'success': True}
         responses.add(
             responses.POST,
@@ -56,10 +56,23 @@ def test_impersonate(self):
         )
         client = Client.from_request(request(), {})
         options = {'properties': {'impersonator': 'admin'}, 'user_id': '1234'}
-        self.assertEqual(client.impersonate(options), response_text)
+        self.assertEqual(client.start_impersonation(options), response_text)
 
     @responses.activate
-    def test_impersonate_failed(self):
+    def test_end_impersonation(self):
+        response_text = {'success': True}
+        responses.add(
+            responses.DELETE,
+            'https://api.castle.io/v1/impersonate',
+            json=response_text,
+            status=200
+        )
+        client = Client.from_request(request(), {})
+        options = {'properties': {'impersonator': 'admin'}, 'user_id': '1234'}
+        self.assertEqual(client.end_impersonation(options), response_text)
+
+    @responses.activate
+    def test_start_impersonation_failed(self):
         response_text = {}
         responses.add(
             responses.POST,
@@ -70,7 +83,21 @@ def test_impersonate_failed(self):
         client = Client.from_request(request(), {})
         options = {'properties': {'impersonator': 'admin'}, 'user_id': '1234'}
         with self.assertRaises(ImpersonationFailed):
-            client.impersonate(options)
+            client.start_impersonation(options)
+
+    @responses.activate
+    def test_end_impersonation_failed(self):
+        response_text = {}
+        responses.add(
+            responses.DELETE,
+            'https://api.castle.io/v1/impersonate',
+            json=response_text,
+            status=200
+        )
+        client = Client.from_request(request(), {})
+        options = {'properties': {'impersonator': 'admin'}, 'user_id': '1234'}
+        with self.assertRaises(ImpersonationFailed):
+            client.end_impersonation(options)
 
     @responses.activate
     def test_identify_tracked_true(self):
diff --git a/castle/test/commands/approve_device_test.py b/castle/test/commands/approve_device_test.py
@@ -9,12 +9,12 @@ def device_token():
 
 
 class CommandsApproveDeviceTestCase(unittest.TestCase):
-    def test_build_no_device_token(self):
+    def test_call_no_device_token(self):
         with self.assertRaises(InvalidParametersError):
-            CommandsApproveDevice.build('')
+            CommandsApproveDevice.call('')
 
-    def test_build(self):
-        command = CommandsApproveDevice.build(device_token())
+    def test_call(self):
+        command = CommandsApproveDevice.call(device_token())
         self.assertIsInstance(command, Command)
         self.assertEqual(command.method, 'put')
         self.assertEqual(command.path, 'devices/1234/approve')
diff --git a/castle/test/commands/authenticate_test.py b/castle/test/commands/authenticate_test.py
@@ -41,7 +41,7 @@ def test_init(self):
         obj = CommandsAuthenticate(context)
         self.assertEqual(obj.context, context)
 
-    def test_build(self):
+    def test_call(self):
         context = {'test': '1'}
         options = default_options_plus(context={'spam': True})
 
@@ -51,46 +51,46 @@ def test_build(self):
         expected = default_command_with_data(**expected_data)
 
         self.assertEqual(CommandsAuthenticate(
-            context).build(options), expected)
+            context).call(options), expected)
 
-    def test_build_no_event(self):
+    def test_call_no_event(self):
         context = {}
         options = default_options()
         options.pop('event')
 
         with self.assertRaises(InvalidParametersError):
-            CommandsAuthenticate(context).build(options)
+            CommandsAuthenticate(context).call(options)
 
-    def test_build_no_user_id(self):
+    def test_call_no_user_id(self):
         context = {}
         options = default_options()
         options.pop('user_id')
 
         expected = default_command_with_data(**options)
 
-        self.assertEqual(CommandsAuthenticate(context).build(options), expected)
+        self.assertEqual(CommandsAuthenticate(context).call(options), expected)
 
-    def test_build_properties_allowed(self):
+    def test_call_properties_allowed(self):
         context = {}
         options = default_options_plus(properties={'test': '1'})
         options.update({'context': context})
 
         expected = default_command_with_data(**options)
 
         self.assertEqual(CommandsAuthenticate(
-            context).build(options), expected)
+            context).call(options), expected)
 
-    def test_build_user_traits_allowed(self):
+    def test_call_user_traits_allowed(self):
         context = {}
         options = default_options_plus(user_traits={'email': 'a@b.com'})
         options.update({'context': context})
 
         expected = default_command_with_data(**options)
 
         self.assertEqual(CommandsAuthenticate(
-            context).build(options), expected)
+            context).call(options), expected)
 
-    def test_build_risk_policy_allowed(self):
+    def test_call_risk_policy_allowed(self):
         context = {}
         risk_policy = {
             'id': 'q-rbeMzBTdW2Fd09sbz55A',
@@ -102,4 +102,4 @@ def test_build_risk_policy_allowed(self):
         options = default_options_plus(risk_policy=risk_policy)
         expected = default_command_with_data(**options)
         self.assertEqual(CommandsAuthenticate(
-            context).build(options), expected)
+            context).call(options), expected)
diff --git a/castle/test/commands/end_impersonation_test.py b/castle/test/commands/end_impersonation_test.py
diff --git a/castle/test/commands/get_device_test.py b/castle/test/commands/get_device_test.py
diff --git a/castle/test/commands/get_devices_for_user_test.py b/castle/test/commands/get_devices_for_user_test.py
diff --git a/castle/test/commands/identify_test.py b/castle/test/commands/identify_test.py
diff --git a/castle/test/commands/report_device_test.py b/castle/test/commands/report_device_test.py
diff --git a/castle/test/commands/review_test.py b/castle/test/commands/review_test.py
diff --git a/castle/test/commands/start_impersonation_test.py b/castle/test/commands/start_impersonation_test.py
diff --git a/castle/test/commands/track_test.py b/castle/test/commands/track_test.py
