Allow for running a script in an arbitrary frame context

Author: pmeenan

internal/devtools.py

@@ -88,6 +88,8 @@ def __init__(self, options, job, task, use_devtools_video, is_webkit, is_ios):
         self.main_thread_blocked = False
         self.stylesheets = {}
         self.headers = {}
+        self.execution_contexts = {}
+        self.execution_context = None
         self.trace_parser = None
         self.prepare()
         self.html_body = False
@@ -564,6 +566,12 @@ def stop_recording(self):
         # Add the audit issues to the page data
         if len(self.audit_issues):
             self.task['page_data']['audit_issues'] = self.audit_issues
+        # Add the list of execution contexts
+        contexts = []
+        for id in self.execution_contexts:
+            contexts.append(self.execution_contexts[id])
+        if len(contexts):
+            self.task['page_data']['execution_contexts'] = contexts
         # Process the timeline data
         if self.trace_parser is not None:
             start = monotonic()
@@ -1174,7 +1182,7 @@ def colors_are_similar(self, color1, color2, threshold=15):
             similar = False
         return similar
 
-    def execute_js(self, script):
+    def execute_js(self, script, use_execution_context=False):
         """Run the provided JS in the browser and return the result"""
         if self.must_exit:
             return
@@ -1183,18 +1191,35 @@ def execute_js(self, script):
             if self.is_webkit:
                 response = self.send_command('Runtime.evaluate', {'expression': script, 'returnByValue': True}, timeout=30, wait=True)
             else:
-                response = self.send_command("Runtime.evaluate",
-                                            {'expression': script,
-                                            'awaitPromise': True,
-                                            'returnByValue': True,
-                                            'timeout': 30000},
-                                            wait=True, timeout=30)
+                params = {'expression': script,
+                          'awaitPromise': True,
+                          'returnByValue': True,
+                          'timeout': 30000}
+                if use_execution_context and self.execution_context is not None:
+                    params['contextId'] = self.execution_context
+                response = self.send_command("Runtime.evaluate", params, wait=True, timeout=30)
             if response is not None and 'result' in response and\
                     'result' in response['result'] and\
                     'value' in response['result']['result']:
                 ret = response['result']['result']['value']
         return ret
 
+    def set_execution_context(self, target):
+        """ Set the js execution context by matching id, origin or name """
+        if len(target):
+            parts = target.split('=', 1)
+            if len(parts) == 2:
+                key = parts[0].strip()
+                value = parts[1].strip()
+                if key in ['id', 'name', 'origin'] and len(value):
+                    for id in self.execution_contexts:
+                        context = self.execution_contexts[id]
+                        if key in context and context[key] == value:
+                            self.execution_context = id
+                            break
+        else:
+            self.execution_context = None
+
     def set_header(self, header):
         """Add/modify a header on the outbound requests"""
         if header is not None and len(header):
@@ -1309,6 +1334,7 @@ def enable_target(self, target_id=None):
             self.send_command('Network.enable', {}, target_id=target_id)
             self.send_command('Console.enable', {}, target_id=target_id)
             self.send_command('Log.enable', {}, target_id=target_id)
+            self.send_command('Runtime.enable', {}, target_id=target_id)
             self.send_command('Log.startViolationsReport', {'config': [{'name': 'discouragedAPIUse', 'threshold': -1}]}, target_id=target_id)
             self.send_command('Audits.enable', {}, target_id=target_id)
             self.job['shaper'].apply(target_id=target_id)
@@ -1423,6 +1449,8 @@ def process_message(self, msg, target_id=None):
                     self.process_css_event(event, msg)
                 elif category == 'Debugger':
                     self.process_debugger_event(event, msg)
+                elif category == 'Runtime':
+                    self.process_runtime_event(event, msg)
                 elif category == 'Target':
                     log_event = False
                     self.process_target_event(event, msg)
@@ -1545,6 +1573,26 @@ def process_debugger_event(self, event, msg):
         if event == 'paused':
             self.send_command('Debugger.resume', {})
 
+    def process_runtime_event(self, event, msg):
+        """Handle Runtime.* events"""
+        if event == 'executionContextCreated':
+            if 'params' in msg and 'context' in msg['params'] and 'id' in msg['params']['context'] and 'origin':
+                context = msg['params']['context']
+                id = context['id']
+                ctx = {'id': id}
+                if 'origin' in context:
+                    ctx['origin'] = context['origin']
+                if 'name' in context:
+                    ctx['name'] = context['name']
+                self.execution_contexts[id] = ctx
+                logging.debug('Execution context created: %s', json.dumps(context))
+        elif event == 'executionContextDestroyed':
+            if 'params' in msg and 'executionContextId' in msg['params']:
+                id = msg['params']['executionContextId']
+                if id in self.execution_contexts:
+                    del self.execution_contexts[id]
+                    logging.debug('Execution context %d deleted', id)
+
     def process_network_event(self, event, msg, target_id=None):
         """Process Network.* dev tools events"""
         if event == 'requestIntercepted':

internal/devtools_browser.py

@@ -689,9 +689,16 @@ def process_command(self, command):
                 needs_mark = True
                 if self.task['combine_steps']:
                     needs_mark = False
-                script = self.prepare_script_for_record(script, needs_mark) #pylint: disable=no-member
+                if self.devtools.execution_context is not None:
+                    # Clear the orange frame as a separate step to make sure it is done in the correct context
+                    clear_script = self.prepare_script_for_record('', needs_mark)
+                    self.devtools.execute_js(clear_script)
+                else:
+                    script = self.prepare_script_for_record(script, needs_mark) #pylint: disable=no-member
                 self.devtools.start_navigating()
-            self.devtools.execute_js(script)
+            self.devtools.execute_js(script, True)
+        elif command['command'] == 'setexecutioncontext':
+            self.devtools.set_execution_context(command['target'])
         elif command['command'] == 'sleep':
             available_sleep = 60 - self.total_sleep
             delay = min(available_sleep, max(0, int(re.search(r'\d+', str(command['target'])).group())))