New features: screenshot, open URLs in new window, query for +windowFocused

brotab/api.py

@@ -162,6 +162,9 @@ def activate_tab(self, args: List[str], focused: bool):
     def get_active_tabs(self, args) -> List[str]:
         return [self.prefix_tab(tab) for tab in self._get('/get_active_tabs').split(',')]
 
+    def get_screenshot(self, args):
+        return self._get('/get_screenshot')
+
     def query_tabs(self, args):
         query = args
         if isinstance(query, str):

brotab/extension/chrome/background.js

@@ -50,6 +50,10 @@ class BrowserTabs {
     throw new Error('getActive is not implemented');
   }
 
+  getActiveScreenshot(onSuccess) {
+    throw new Error('getActiveScreenshot is not implemented');
+  }
+
   runScript(tab_id, script, payload, onSuccess, onError) {
     throw new Error('runScript is not implemented');
   }
@@ -68,10 +72,30 @@ class FirefoxTabs extends BrowserTabs {
   }
 
   query(queryInfo, onSuccess) {
-    this._browser.tabs.query(queryInfo).then(
-      onSuccess,
-      (error) => console.log(`Error executing queryTabs: ${error}`)
-    );
+    if (queryInfo.hasOwnProperty('windowFocused')) {
+      let keepFocused = queryInfo['windowFocused']
+      delete queryInfo.windowFocused;
+      this._browser.tabs.query(queryInfo).then(
+        tabs => {
+          Promise.all(tabs.map(tab => {
+            return new Promise(resolve => {
+              this._browser.windows.get(tab.windowId, {populate: false}, window => {
+                resolve(window.focused === keepFocused ? tab : null);
+              });
+            });
+          })).then(result => {
+            tabs = result.filter(tab => tab !== null);
+            onSuccess(tabs);
+          });
+        },
+        (error) => console.log(`Error executing queryTabs: ${error}`)
+      );
+    } else {
+      this._browser.tabs.query(queryInfo).then(
+        onSuccess,
+        (error) => console.log(`Error executing queryTabs: ${error}`)
+      );
+    }
   }
 
   close(tab_ids, onSuccess) {
@@ -100,10 +124,17 @@ class FirefoxTabs extends BrowserTabs {
   }
 
   create(createOptions, onSuccess) {
-    this._browser.tabs.create(createOptions).then(
-      onSuccess,
-      (error) => console.log(`Error: ${error}`)
-    );
+    if (createOptions.windowId === 0) {
+      this._browser.windows.create({ url: createOptions.url }).then(
+        onSuccess,
+        (error) => console.log(`Error: ${error}`)
+      );
+    } else {
+      this._browser.tabs.create(createOptions).then(
+        onSuccess,
+        (error) => console.log(`Error: ${error}`)
+      );
+    }
   }
 
   getActive(onSuccess) {
@@ -113,6 +144,29 @@ class FirefoxTabs extends BrowserTabs {
     );
   }
 
+  getActiveScreenshot(onSuccess) {
+    let queryOptions = { active: true, lastFocusedWindow: true };
+    this._browser.tabs.query(queryOptions).then(
+      (tabs) => {
+        let tab = tabs[0];
+        let windowId = tab.windowId;
+        let tabId = tab.id;
+        this._browser.tabs.captureVisibleTab(windowId, { format: 'png' }).then(
+          function(data) {
+            const message = {
+              tab: tabId,
+              window: windowId,
+              data: data
+            };
+            onSuccess(message);
+          },
+          (error) => console.log(`Error: ${error}`)
+        );
+      },
+      (error) => console.log(`Error: ${error}`)
+    );
+  }
+
   runScript(tab_id, script, payload, onSuccess, onError) {
     this._browser.tabs.executeScript(tab_id, {code: script}).then(
       (result) => onSuccess(result, payload),
@@ -145,7 +199,24 @@ class ChromeTabs extends BrowserTabs {
   }
 
   query(queryInfo, onSuccess) {
-    this._browser.tabs.query(queryInfo, onSuccess);
+    if (queryInfo.hasOwnProperty('windowFocused')) {
+      let keepFocused = queryInfo['windowFocused']
+      delete queryInfo.windowFocused;
+      this._browser.tabs.query(queryInfo, tabs => {
+        Promise.all(tabs.map(tab => {
+          return new Promise(resolve => {
+            this._browser.windows.get(tab.windowId, {populate: false}, window => {
+              resolve(window.focused === keepFocused ? tab : null);
+            });
+          });
+        })).then(result => {
+          tabs = result.filter(tab => tab !== null);
+          onSuccess(tabs);
+        });
+      });
+    } else {
+      this._browser.tabs.query(queryInfo, onSuccess);
+    }
   }
 
   close(tab_ids, onSuccess) {
@@ -169,13 +240,35 @@ class ChromeTabs extends BrowserTabs {
   }
 
   create(createOptions, onSuccess) {
-    this._browser.tabs.create(createOptions, onSuccess);
+    if (createOptions.windowId === 0) {
+      this._browser.windows.create({ url: createOptions.url }, onSuccess);
+    } else {
+      this._browser.tabs.create(createOptions, onSuccess);
+    }
   }
 
   getActive(onSuccess) {
     this._browser.tabs.query({active: true}, onSuccess);
   }
 
+  getActiveScreenshot(onSuccess) {
+    // this._browser.tabs.captureVisibleTab(null, { format: 'png' }, onSuccess);
+    let queryOptions = { active: true, lastFocusedWindow: true };
+    this._browser.tabs.query(queryOptions, (tabs) => {
+      let tab = tabs[0];
+      let windowId = tab.windowId;
+      let tabId = tab.id;
+      this._browser.tabs.captureVisibleTab(windowId, { format: 'png' }, function(data) {
+        const message = {
+          tab: tabId,
+          window: windowId,
+          data: data
+        };
+        onSuccess(message);
+      });
+    });
+  }
+
   runScript(tab_id, script, payload, onSuccess, onError) {
     this._browser.tabs.executeScript(
       tab_id, {code: script},
@@ -282,7 +375,7 @@ function queryTabs(query_info) {
 
     integerKeys = {'windowId': null, 'index': null};
     booleanKeys = {'active': null, 'pinned': null, 'audible': null, 'muted': null, 'highlighted': null,
-      'discarded': null, 'autoDiscardable': null, 'currentWindow': null, 'lastFocusedWindow': null};
+      'discarded': null, 'autoDiscardable': null, 'currentWindow': null, 'lastFocusedWindow': null, 'windowFocused': null};
 
     query = Object.entries(query).reduce((o, [k,v]) => {
       if (booleanKeys.hasOwnProperty(k) && typeof v != 'boolean') {
@@ -335,13 +428,23 @@ function closeTabs(tab_ids) {
   browserTabs.close(tab_ids, () => port.postMessage('OK'));
 }
 
-function openUrls(urls, window_id) {
+function openUrls(urls, window_id, first_result="") {
   if (urls.length == 0) {
     console.log('Opening urls done');
     port.postMessage([]);
     return;
   }
 
+  if (window_id === 0) {
+    browserTabs.create({'url': urls[0], windowId: 0}, (window) => {
+      result = `${window.id}.${window.tabs[0].id}`;
+      console.log(`Opened first window: ${result}`);
+      urls = urls.slice(1);
+      openUrls(urls, window.id, result);
+    });
+    return;
+  }
+
   var promises = [];
   for (let url of urls) {
     console.log(`Opening another one url ${url}`);
@@ -352,6 +455,9 @@ function openUrls(urls, window_id) {
     }))
   };
   Promise.all(promises).then(result => {
+    if (first_result !== "") {
+      result.unshift(first_result);
+    }
     const data = Array.prototype.concat(...result)
     console.log(`Sending ids back: ${JSON.stringify(data)}`);
     port.postMessage(data)
@@ -405,6 +511,12 @@ function getActiveTabs() {
   });
 }
 
+function getActiveScreenshot() {
+  browserTabs.getActiveScreenshot(data => {
+    port.postMessage(data);
+  });
+}
+
 function getWordsScript(match_regex, join_with) {
   return GET_WORDS_SCRIPT
     .replace('#match_regex#', match_regex)
@@ -609,6 +721,11 @@ port.onMessage.addListener((command) => {
     getActiveTabs();
   }
 
+  else if (command['name'] == 'get_screenshot') {
+    console.log('Getting visible screenshot');
+    getActiveScreenshot();
+  }
+
   else if (command['name'] == 'get_words') {
     console.log('Getting words from tab:', command['tab_id']);
     getWords(command['tab_id'], command['match_regex'], command['join_with']);

brotab/main.py

@@ -54,7 +54,7 @@
 from argparse import ArgumentParser
 from functools import partial
 from itertools import groupby
-from json import loads
+from json import loads, dumps
 from string import ascii_lowercase
 from typing import List
 from typing import Tuple
@@ -173,6 +173,17 @@ def show_active_tabs(args):
             print('%s\t%s' % (tab, api))
 
 
+def screenshot(args):
+    brotab_logger.info('Getting screenshot: %s', args)
+    apis = create_clients(args.target_hosts)
+    for api in apis:
+        result = api.get_screenshot(args)
+        # print(result, api)
+        result = loads(result)
+        result['api'] = api._prefix[:1]
+        result = dumps(result)
+        print(result)
+
 def search_tabs(args):
     for result in query(args.sqlite, args.query):
         print('\t'.join([result.tab_id, result.title, result.snippet]))
@@ -485,6 +496,13 @@ def parse_args(args):
         ''')
     parser_active_tab.set_defaults(func=show_active_tabs)
 
+    parser_screenshot = subparsers.add_parser(
+        'screenshot',
+        help='''
+        return base64 screenshot in json object with keys: 'data' (base64 png), 'tab' (tab id of visible tab), 'window' (window id of visible tab), 'api' (prefix of client api)
+        ''')
+    parser_screenshot.set_defaults(func=screenshot)
+
     parser_search_tabs = subparsers.add_parser(
         'search',
         help='''
@@ -536,6 +554,10 @@ def parse_args(args):
                                    help='tabs are in the last focused window.')
     parser_query_tabs.add_argument('-lastFocusedWindow', action='store_const', const=False, default=None,
                                    help='tabs are not in the last focused window.')
+    parser_query_tabs.add_argument('+windowFocused', action='store_const', const=True, default=None,
+                                   help='tabs are in the focused window.')
+    parser_query_tabs.add_argument('-windowFocused', action='store_const', const=False, default=None,
+                                   help='tabs are not in the focused window.')
     parser_query_tabs.add_argument('-status', type=str, choices=['loading', 'complete'],
                                    help='whether the tabs have completed loading i.e. loading or complete.')
     parser_query_tabs.add_argument('-title', type=str,
@@ -594,7 +616,7 @@ def parse_args(args):
         open URLs from the stdin (one URL per line). One positional argument is
         required: <prefix>.<window_id> OR <client>. If window_id is not
         specified, URL will be opened in the active window of the specifed
-        client
+        client. If window_id is 0, URLs will be opened in new window.
         ''')
     parser_open_urls.set_defaults(func=open_urls)
     parser_open_urls.add_argument(

brotab/mediator/http_server.py

@@ -67,6 +67,7 @@ def _setup_routes(self) -> None:
         self.app.route('/new_tab/<query>', methods=['GET'])(self.new_tab)
         self.app.route('/activate_tab/<int:tab_id>', methods=['GET'])(self.activate_tab)
         self.app.route('/get_active_tabs', methods=['GET'])(self.get_active_tabs)
+        self.app.route('/get_screenshot', methods=['GET'])(self.get_screenshot)
         self.app.route('/get_words/<string:tab_id>', methods=['GET'])(self.get_words)
         self.app.route('/get_words', methods=['GET'])(self.get_words)
         self.app.route('/get_text', methods=['GET'])(self.get_text)
@@ -141,6 +142,9 @@ def activate_tab(self, tab_id):
     def get_active_tabs(self):
         return self.remote_api.get_active_tabs()
 
+    def get_screenshot(self):
+        return self.remote_api.get_screenshot()
+
     def get_words(self, tab_id=None):
         tab_id = int(tab_id) if is_valid_integer(tab_id) else None
         match_regex = request.args.get('match_regex', DEFAULT_GET_WORDS_MATCH_REGEX)

brotab/mediator/remote_api.py

@@ -97,6 +97,12 @@ def get_active_tabs(self) -> str:
         self._transport.send(command)
         return self._transport.recv()
 
+    def get_screenshot(self) -> str:
+        mediator_logger.info('getting screemsjpt')
+        command = {'name': 'get_screenshot'}
+        self._transport.send(command)
+        return self._transport.recv()
+
     def get_words(self, tab_id: str, match_regex: str, join_with: str):
         mediator_logger.info('getting tab words: %s', tab_id)
         command = {

requirements/base.txt

@@ -1,3 +1,4 @@
 Flask==2.0.2
 requests==2.24.0
-psutil==5.8.0
+psutil==5.8.0
+Werkzeug==2.2.2