Skip to content

Commit a51ecb9

Browse files
balta2arbalta2ar
committed
Merge branch 'digitalsignalperson-2023-10-features'
2 parents 02e76a5 + e60e373 commit a51ecb9

File tree

6 files changed

+167
-15
lines changed

6 files changed

+167
-15
lines changed

‎brotab/api.py

+3
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,9 @@ def activate_tab(self, args: List[str], focused: bool):
162162
def get_active_tabs(self, args) -> List[str]:
163163
return [self.prefix_tab(tab) for tab in self._get('/get_active_tabs').split(',')]
164164

165+
def get_screenshot(self, args):
166+
return self._get('/get_screenshot')
167+
165168
def query_tabs(self, args):
166169
query = args
167170
if isinstance(query, str):

‎brotab/extension/chrome/background.js

+129-12
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,10 @@ class BrowserTabs {
5050
throw new Error('getActive is not implemented');
5151
}
5252

53+
getActiveScreenshot(onSuccess) {
54+
throw new Error('getActiveScreenshot is not implemented');
55+
}
56+
5357
runScript(tab_id, script, payload, onSuccess, onError) {
5458
throw new Error('runScript is not implemented');
5559
}
@@ -68,10 +72,30 @@ class FirefoxTabs extends BrowserTabs {
6872
}
6973

7074
query(queryInfo, onSuccess) {
71-
this._browser.tabs.query(queryInfo).then(
72-
onSuccess,
73-
(error) => console.log(`Error executing queryTabs: ${error}`)
74-
);
75+
if (queryInfo.hasOwnProperty('windowFocused')) {
76+
let keepFocused = queryInfo['windowFocused']
77+
delete queryInfo.windowFocused;
78+
this._browser.tabs.query(queryInfo).then(
79+
tabs => {
80+
Promise.all(tabs.map(tab => {
81+
return new Promise(resolve => {
82+
this._browser.windows.get(tab.windowId, {populate: false}, window => {
83+
resolve(window.focused === keepFocused ? tab : null);
84+
});
85+
});
86+
})).then(result => {
87+
tabs = result.filter(tab => tab !== null);
88+
onSuccess(tabs);
89+
});
90+
},
91+
(error) => console.log(`Error executing queryTabs: ${error}`)
92+
);
93+
} else {
94+
this._browser.tabs.query(queryInfo).then(
95+
onSuccess,
96+
(error) => console.log(`Error executing queryTabs: ${error}`)
97+
);
98+
}
7599
}
76100

77101
close(tab_ids, onSuccess) {
@@ -100,10 +124,17 @@ class FirefoxTabs extends BrowserTabs {
100124
}
101125

102126
create(createOptions, onSuccess) {
103-
this._browser.tabs.create(createOptions).then(
104-
onSuccess,
105-
(error) => console.log(`Error: ${error}`)
106-
);
127+
if (createOptions.windowId === 0) {
128+
this._browser.windows.create({ url: createOptions.url }).then(
129+
onSuccess,
130+
(error) => console.log(`Error: ${error}`)
131+
);
132+
} else {
133+
this._browser.tabs.create(createOptions).then(
134+
onSuccess,
135+
(error) => console.log(`Error: ${error}`)
136+
);
137+
}
107138
}
108139

109140
getActive(onSuccess) {
@@ -113,6 +144,29 @@ class FirefoxTabs extends BrowserTabs {
113144
);
114145
}
115146

147+
getActiveScreenshot(onSuccess) {
148+
let queryOptions = { active: true, lastFocusedWindow: true };
149+
this._browser.tabs.query(queryOptions).then(
150+
(tabs) => {
151+
let tab = tabs[0];
152+
let windowId = tab.windowId;
153+
let tabId = tab.id;
154+
this._browser.tabs.captureVisibleTab(windowId, { format: 'png' }).then(
155+
function(data) {
156+
const message = {
157+
tab: tabId,
158+
window: windowId,
159+
data: data
160+
};
161+
onSuccess(message);
162+
},
163+
(error) => console.log(`Error: ${error}`)
164+
);
165+
},
166+
(error) => console.log(`Error: ${error}`)
167+
);
168+
}
169+
116170
runScript(tab_id, script, payload, onSuccess, onError) {
117171
this._browser.tabs.executeScript(tab_id, {code: script}).then(
118172
(result) => onSuccess(result, payload),
@@ -145,7 +199,24 @@ class ChromeTabs extends BrowserTabs {
145199
}
146200

147201
query(queryInfo, onSuccess) {
148-
this._browser.tabs.query(queryInfo, onSuccess);
202+
if (queryInfo.hasOwnProperty('windowFocused')) {
203+
let keepFocused = queryInfo['windowFocused']
204+
delete queryInfo.windowFocused;
205+
this._browser.tabs.query(queryInfo, tabs => {
206+
Promise.all(tabs.map(tab => {
207+
return new Promise(resolve => {
208+
this._browser.windows.get(tab.windowId, {populate: false}, window => {
209+
resolve(window.focused === keepFocused ? tab : null);
210+
});
211+
});
212+
})).then(result => {
213+
tabs = result.filter(tab => tab !== null);
214+
onSuccess(tabs);
215+
});
216+
});
217+
} else {
218+
this._browser.tabs.query(queryInfo, onSuccess);
219+
}
149220
}
150221

151222
close(tab_ids, onSuccess) {
@@ -169,13 +240,35 @@ class ChromeTabs extends BrowserTabs {
169240
}
170241

171242
create(createOptions, onSuccess) {
172-
this._browser.tabs.create(createOptions, onSuccess);
243+
if (createOptions.windowId === 0) {
244+
this._browser.windows.create({ url: createOptions.url }, onSuccess);
245+
} else {
246+
this._browser.tabs.create(createOptions, onSuccess);
247+
}
173248
}
174249

175250
getActive(onSuccess) {
176251
this._browser.tabs.query({active: true}, onSuccess);
177252
}
178253

254+
getActiveScreenshot(onSuccess) {
255+
// this._browser.tabs.captureVisibleTab(null, { format: 'png' }, onSuccess);
256+
let queryOptions = { active: true, lastFocusedWindow: true };
257+
this._browser.tabs.query(queryOptions, (tabs) => {
258+
let tab = tabs[0];
259+
let windowId = tab.windowId;
260+
let tabId = tab.id;
261+
this._browser.tabs.captureVisibleTab(windowId, { format: 'png' }, function(data) {
262+
const message = {
263+
tab: tabId,
264+
window: windowId,
265+
data: data
266+
};
267+
onSuccess(message);
268+
});
269+
});
270+
}
271+
179272
runScript(tab_id, script, payload, onSuccess, onError) {
180273
this._browser.tabs.executeScript(
181274
tab_id, {code: script},
@@ -282,7 +375,7 @@ function queryTabs(query_info) {
282375

283376
integerKeys = {'windowId': null, 'index': null};
284377
booleanKeys = {'active': null, 'pinned': null, 'audible': null, 'muted': null, 'highlighted': null,
285-
'discarded': null, 'autoDiscardable': null, 'currentWindow': null, 'lastFocusedWindow': null};
378+
'discarded': null, 'autoDiscardable': null, 'currentWindow': null, 'lastFocusedWindow': null, 'windowFocused': null};
286379

287380
query = Object.entries(query).reduce((o, [k,v]) => {
288381
if (booleanKeys.hasOwnProperty(k) && typeof v != 'boolean') {
@@ -335,13 +428,23 @@ function closeTabs(tab_ids) {
335428
browserTabs.close(tab_ids, () => port.postMessage('OK'));
336429
}
337430

338-
function openUrls(urls, window_id) {
431+
function openUrls(urls, window_id, first_result="") {
339432
if (urls.length == 0) {
340433
console.log('Opening urls done');
341434
port.postMessage([]);
342435
return;
343436
}
344437

438+
if (window_id === 0) {
439+
browserTabs.create({'url': urls[0], windowId: 0}, (window) => {
440+
result = `${window.id}.${window.tabs[0].id}`;
441+
console.log(`Opened first window: ${result}`);
442+
urls = urls.slice(1);
443+
openUrls(urls, window.id, result);
444+
});
445+
return;
446+
}
447+
345448
var promises = [];
346449
for (let url of urls) {
347450
console.log(`Opening another one url ${url}`);
@@ -352,6 +455,9 @@ function openUrls(urls, window_id) {
352455
}))
353456
};
354457
Promise.all(promises).then(result => {
458+
if (first_result !== "") {
459+
result.unshift(first_result);
460+
}
355461
const data = Array.prototype.concat(...result)
356462
console.log(`Sending ids back: ${JSON.stringify(data)}`);
357463
port.postMessage(data)
@@ -405,6 +511,12 @@ function getActiveTabs() {
405511
});
406512
}
407513

514+
function getActiveScreenshot() {
515+
browserTabs.getActiveScreenshot(data => {
516+
port.postMessage(data);
517+
});
518+
}
519+
408520
function getWordsScript(match_regex, join_with) {
409521
return GET_WORDS_SCRIPT
410522
.replace('#match_regex#', match_regex)
@@ -609,6 +721,11 @@ port.onMessage.addListener((command) => {
609721
getActiveTabs();
610722
}
611723

724+
else if (command['name'] == 'get_screenshot') {
725+
console.log('Getting visible screenshot');
726+
getActiveScreenshot();
727+
}
728+
612729
else if (command['name'] == 'get_words') {
613730
console.log('Getting words from tab:', command['tab_id']);
614731
getWords(command['tab_id'], command['match_regex'], command['join_with']);

‎brotab/main.py

+24-2
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@
5454
from argparse import ArgumentParser
5555
from functools import partial
5656
from itertools import groupby
57-
from json import loads
57+
from json import loads, dumps
5858
from string import ascii_lowercase
5959
from typing import List
6060
from typing import Tuple
@@ -173,6 +173,17 @@ def show_active_tabs(args):
173173
print('%s\t%s' % (tab, api))
174174

175175

176+
def screenshot(args):
177+
brotab_logger.info('Getting screenshot: %s', args)
178+
apis = create_clients(args.target_hosts)
179+
for api in apis:
180+
result = api.get_screenshot(args)
181+
# print(result, api)
182+
result = loads(result)
183+
result['api'] = api._prefix[:1]
184+
result = dumps(result)
185+
print(result)
186+
176187
def search_tabs(args):
177188
for result in query(args.sqlite, args.query):
178189
print('\t'.join([result.tab_id, result.title, result.snippet]))
@@ -485,6 +496,13 @@ def parse_args(args):
485496
''')
486497
parser_active_tab.set_defaults(func=show_active_tabs)
487498

499+
parser_screenshot = subparsers.add_parser(
500+
'screenshot',
501+
help='''
502+
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)
503+
''')
504+
parser_screenshot.set_defaults(func=screenshot)
505+
488506
parser_search_tabs = subparsers.add_parser(
489507
'search',
490508
help='''
@@ -536,6 +554,10 @@ def parse_args(args):
536554
help='tabs are in the last focused window.')
537555
parser_query_tabs.add_argument('-lastFocusedWindow', action='store_const', const=False, default=None,
538556
help='tabs are not in the last focused window.')
557+
parser_query_tabs.add_argument('+windowFocused', action='store_const', const=True, default=None,
558+
help='tabs are in the focused window.')
559+
parser_query_tabs.add_argument('-windowFocused', action='store_const', const=False, default=None,
560+
help='tabs are not in the focused window.')
539561
parser_query_tabs.add_argument('-status', type=str, choices=['loading', 'complete'],
540562
help='whether the tabs have completed loading i.e. loading or complete.')
541563
parser_query_tabs.add_argument('-title', type=str,
@@ -594,7 +616,7 @@ def parse_args(args):
594616
open URLs from the stdin (one URL per line). One positional argument is
595617
required: <prefix>.<window_id> OR <client>. If window_id is not
596618
specified, URL will be opened in the active window of the specifed
597-
client
619+
client. If window_id is 0, URLs will be opened in new window.
598620
''')
599621
parser_open_urls.set_defaults(func=open_urls)
600622
parser_open_urls.add_argument(

‎brotab/mediator/http_server.py

+4
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ def _setup_routes(self) -> None:
6767
self.app.route('/new_tab/<query>', methods=['GET'])(self.new_tab)
6868
self.app.route('/activate_tab/<int:tab_id>', methods=['GET'])(self.activate_tab)
6969
self.app.route('/get_active_tabs', methods=['GET'])(self.get_active_tabs)
70+
self.app.route('/get_screenshot', methods=['GET'])(self.get_screenshot)
7071
self.app.route('/get_words/<string:tab_id>', methods=['GET'])(self.get_words)
7172
self.app.route('/get_words', methods=['GET'])(self.get_words)
7273
self.app.route('/get_text', methods=['GET'])(self.get_text)
@@ -141,6 +142,9 @@ def activate_tab(self, tab_id):
141142
def get_active_tabs(self):
142143
return self.remote_api.get_active_tabs()
143144

145+
def get_screenshot(self):
146+
return self.remote_api.get_screenshot()
147+
144148
def get_words(self, tab_id=None):
145149
tab_id = int(tab_id) if is_valid_integer(tab_id) else None
146150
match_regex = request.args.get('match_regex', DEFAULT_GET_WORDS_MATCH_REGEX)

‎brotab/mediator/remote_api.py

+6
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,12 @@ def get_active_tabs(self) -> str:
9797
self._transport.send(command)
9898
return self._transport.recv()
9999

100+
def get_screenshot(self) -> str:
101+
mediator_logger.info('getting screemsjpt')
102+
command = {'name': 'get_screenshot'}
103+
self._transport.send(command)
104+
return self._transport.recv()
105+
100106
def get_words(self, tab_id: str, match_regex: str, join_with: str):
101107
mediator_logger.info('getting tab words: %s', tab_id)
102108
command = {

‎requirements/base.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@ Flask==2.2.2
22
requests==2.24.0
33
psutil==5.8.0
44
Werkzeug<3.0
5-
setuptools==75.8.0
5+
setuptools==75.8.0

0 commit comments

Comments
 (0)