tor-browser

The Tor Browser
git clone https://git.dasho.dev/tor-browser.git
Log | Files | Refs | README | LICENSE

decorators.py (2209B)

      1 # This Source Code Form is subject to the terms of the Mozilla Public
      2 # License, v. 2.0. If a copy of the MPL was not distributed with this
      3 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
      4 
      5 import socket
      6 from functools import wraps
      7 
      8 
      9 def _find_marionette_in_args(*args, **kwargs):
     10    try:
     11        m = [a for a in args + tuple(kwargs.values()) if hasattr(a, "session")][0]
     12    except IndexError:
     13        print("Can only apply decorator to function using a marionette object")
     14        raise
     15    return m
     16 
     17 
     18 def do_process_check(func):
     19    """Decorator which checks the process status after the function has run."""
     20 
     21    @wraps(func)
     22    def _(*args, **kwargs):
     23        try:
     24            return func(*args, **kwargs)
     25        except (OSError, socket.timeout):
     26            m = _find_marionette_in_args(*args, **kwargs)
     27 
     28            # In case of socket failures which will also include crashes of the
     29            # application, make sure to handle those correctly. In case of an
     30            # active shutdown just let it bubble up.
     31            if m.is_shutting_down:
     32                raise
     33 
     34            m._handle_socket_failure()
     35 
     36    return _
     37 
     38 
     39 def uses_marionette(func):
     40    """Decorator which creates a marionette session and deletes it
     41    afterwards if one doesn't already exist.
     42    """
     43 
     44    @wraps(func)
     45    def _(*args, **kwargs):
     46        m = _find_marionette_in_args(*args, **kwargs)
     47        delete_session = False
     48        if not m.session:
     49            delete_session = True
     50            m.start_session()
     51 
     52        m.set_context(m.CONTEXT_CHROME)
     53        ret = func(*args, **kwargs)
     54 
     55        if delete_session:
     56            m.delete_session()
     57 
     58        return ret
     59 
     60    return _
     61 
     62 
     63 def using_context(context):
     64    """Decorator which allows a function to execute in certain scope
     65    using marionette.using_context functionality and returns to old
     66    scope once the function exits.
     67    :param context: Either 'chrome' or 'content'.
     68    """
     69 
     70    def wrap(func):
     71        @wraps(func)
     72        def inner(*args, **kwargs):
     73            m = _find_marionette_in_args(*args, **kwargs)
     74            with m.using_context(context):
     75                return func(*args, **kwargs)
     76 
     77        return inner
     78 
     79    return wrap