diff --git a/grattoir/__init__.py b/grattoir/__init__.py index e019c4c..258e3eb 100644 --- a/grattoir/__init__.py +++ b/grattoir/__init__.py @@ -1,60 +1,97 @@ -from selenium.webdriver import firefox from fake_useragent import UserAgent from pathlib import Path +from selenium.webdriver import firefox import importlib import os +import tempfile import types +import uuid -def tor_profile(): - tor_firefox_path = str(Path("~/.tor project/firefox").expanduser()) - ini_path = tor_firefox_path + "/profiles.ini" - with open(ini_path) as ini: - for line in ini: - if "Path=" in line: - return tor_firefox_path + '/' + line.split('=')[1][:-1] - - -def tor_browser_binary_path(): - for pe in os.environ["PATH"].split(os.pathsep): - checkname = os.path.join(pe, "tor-browser") - if os.access(checkname, os.X_OK) and not os.path.isdir(checkname): - return checkname - return None - - -def install_har_export_trigger(driver): - with importlib.resources.path( - __package__ + ".assets", - "har_export_trigger-0.6.2resigned1.xpi", - ) as p: - driver.install_addon(p.as_uri()[7:]) - return driver - - -def har_data(self): - return self.execute_async_script( - "HAR.triggerExport().then(arguments[0]);" -) - +def unique_name(prefix): + return prefix + str(uuid.uuid4()) class Grattoir(object): - def __enter__(self): + def __init__(self, name, with_har=False): + self.name = name + self.with_har = with_har + + def init_const(self, name, with_har): + self.name = name + if self.with_har: + self.har_filename = unique_name(self.name) + self.har_export_path = tempfile.TemporaryDirectory() + self.export_path = ( + str(self.har_export_path) + "/" + self.har_filename + ) + + def tor_browser_binary_path(self): + for pe in os.environ["PATH"].split(os.pathsep): + checkname = os.path.join(pe, "tor-browser") + if os.access(checkname, os.X_OK) and not os.path.isdir(checkname): + return checkname + return None + + def install_har_export_trigger(self): + with importlib.resources.path( + __package__ + ".assets", + "har_export_trigger-0.6.2resigned1.xpi", + ) as p: + self._driver.install_addon(p.as_uri()[7:]) + + def tor_profile_path(self): + tor_firefox_path = str(Path("~/.tor project/firefox").expanduser()) + ini_path = tor_firefox_path + "/profiles.ini" + with open(ini_path) as ini: + for line in ini: + if "Path=" in line: + return tor_firefox_path + '/' + line.split('=')[1][:-1] + + def har_enabled_profile(self, profile): + prefs = { + "extensions.netmonitor.har.enableAutomation": True, + "extensions.netmonitor.har.contentAPIToken": "test", + "extensions.netmonitor.har.autoConnect": True, + "devtools.netmonitor.har.compress": False, + "devtools.netmonitor.har.defaultFileName": self.har_filename, + "devtools.netmonitor.har.defaultLogDir": str(self.har_export_path), + "devtools.netmonitor.har.enableAutoExportToFile": True, + "devtools.netmonitor.har.forceExport": True, + "devtools.netmonitor.har.includeResponseBodies": True, + "devtools.netmonitor.har.jsonp": True, + "devtools.netmonitor.har.jsonpCallback": True, + "devtools.netmonitor.har.pageLoadedTimeout": "2500", + "devtools.toolbox.selectedTool": "netmonitor", + "devtools.netmonitor.persistlog": True, + } + for k, v in prefs.items(): + profile.set_preference(k, v) + return profile + + def firefox_profile(self): + profile = firefox.firefox_profile.FirefoxProfile(profile_directory=self.tor_profile_path()) + if self.with_har: + profile = har_enabled_profile(profile) + return profile + + def init_driver(self): self._driver = None - p_path = tor_profile() - profile = firefox.firefox_profile.FirefoxProfile(profile_directory=p_path) - profile.set_preference("devtools.toolbox.selectedTool", "netmonitor") - profile.set_preference("devtools.netmonitor.persistlog", True) - b_path = tor_browser_binary_path() + b_path = self.tor_browser_binary_path() binary = firefox.firefox_binary.FirefoxBinary(firefox_path=b_path) - o = firefox.options.Options() - o.profile = profile - o.binary = binary - self._driver = firefox.webdriver.WebDriver(options=o) - self._driver = install_har_export_trigger(self._driver) - self._driver.har_data = types.MethodType(har_data, self._driver) + profile = self.firefox_profile() + options = firefox.options.Options() + options.profile = profile + options.binary = binary + self._driver = firefox.webdriver.WebDriver(options=options) + self.install_har_export_trigger() + + def __enter__(self): + self.init_const(self.name, self.with_har) + self.init_driver() return self._driver def __exit__(self, exc_type, exc_value, exc_tb): self._driver.quit() + if self.with_har: + self.har_export_path.cleanup()