Change send_tor_signal arg type, update function doc

send_tor_signal now accepts a stem.Signal arg (a bit cleaner tbh). Also
added the doc string for the "disable" attribute in TorError.
This commit is contained in:
Ben Busby 2020-10-25 21:25:15 -04:00
parent c51dad7529
commit ea88f4bd1d
2 changed files with 9 additions and 8 deletions

View File

@ -4,6 +4,7 @@ from app.utils.gen_ddg_bangs import gen_bangs_json
from flask import Flask
from flask_session import Session
import os
from stem import Signal
app = Flask(__name__, static_folder=os.path.dirname(os.path.abspath(__file__)) + '/static')
app.user_elements = {}
@ -34,6 +35,6 @@ if not os.path.exists(app.config['BANG_PATH']):
Session(app)
# Attempt to acquire tor identity, to determine if Tor config is available
send_tor_signal()
send_tor_signal(Signal.HEARTBEAT)
from app import routes

View File

@ -24,6 +24,9 @@ class TorError(Exception):
Attributes:
message -- a message describing the error that occurred
disable -- optionally disables Tor in the user config (note:
this should only happen if the connection has been dropped
altogether).
"""
def __init__(self, message, disable=False):
@ -32,14 +35,11 @@ class TorError(Exception):
super().__init__(self.message)
def send_tor_signal(new_identity=False) -> bool:
if new_identity:
print('Requesting new identity...')
def send_tor_signal(signal: Signal) -> bool:
try:
with Controller.from_port(port=9051) as c:
c.authenticate()
c.signal(Signal.NEWNYM if new_identity else Signal.HEARTBEAT)
c.signal(signal)
os.environ['TOR_AVAILABLE'] = '1'
return True
except (SocketError, ConnectionRefusedError, ConnectionError):
@ -129,7 +129,7 @@ class Request:
def __init__(self, normal_ua, root_path, config: Config):
# Send heartbeat to Tor, used in determining if the user can or cannot
# enable Tor for future requests
send_tor_signal()
send_tor_signal(Signal.HEARTBEAT)
self.language = config.lang_search
self.mobile = 'Android' in normal_ua or 'iPhone' in normal_ua
@ -195,7 +195,7 @@ class Request:
'User-Agent': self.modified_user_agent
}
if self.tor and not send_tor_signal(new_identity=attempt > 0): # Request new identity if the last one failed
if self.tor and not send_tor_signal(Signal.NEWNYM): # Request new identity if the last one failed
raise TorError("Tor was previously enabled, but the connection has been dropped. Please check your " +
"Tor configuration and try again.", disable=True)