Add ability to set new custom speedtest servers (not mini)

This commit is contained in:
Matt Martz 2018-03-23 11:06:57 -05:00
parent 5d0c62faec
commit 59d4172446
1 changed files with 140 additions and 19 deletions

View File

@ -318,6 +318,16 @@ class InvalidSpeedtestMiniServer(SpeedtestException):
"""
class SpeedtestCustomConnectFailure(SpeedtestException):
"""Could not connect to the provided speedtest custom server"""
class InvalidSpeedtestCustomServer(SpeedtestException):
"""Server provided as a speedtest custom server does not actually appear
to be a speedtest custom server
"""
class ShareResultsConnectFailure(SpeedtestException):
"""Could not connect to speedtest.net API to POST results"""
@ -769,6 +779,8 @@ class SocketTestBase(threading.Thread):
self.size = size
self.remaining = self.size
self._address = address
if shutdown_event:
self._shutdown_event = shutdown_event
else:
@ -1111,16 +1123,17 @@ class Speedtest(object):
self._use_socket = use_socket
self.get_config()
if config is not None:
self.config.update(config)
else:
self.get_config()
self.servers = {}
self.closest = []
self._best = {}
self.results = SpeedtestResults(
client=self.config['client'],
client=self.config.get('client'),
opener=self._opener,
secure=secure,
)
@ -1438,6 +1451,87 @@ class Speedtest(object):
return self.servers
def set_custom_server(self, url, include=None, exclude=None):
request = build_request(url)
uh, e = catch_request(request, opener=self._opener)
if e:
raise SpeedtestCustomConnectFailure(
'Failed to connect to %s' % url
)
else:
text = uh.read()
uh.close()
match = re.search('window.ST_PARAMS = (\{.*\});'.encode(), text)
try:
params = json.loads(match.group(1))
except (TypeError, ValueError):
e = get_exception()
printer('ERROR: %r' % e, debug=True)
raise InvalidSpeedtestCustomServer(
'Invalid Speedtest Custom Server: %s' % url
)
test_globals = params['testGlobals']
config = {
'client': {
'ip': test_globals['ipAddress'],
'lat': test_globals['location']['latitude'],
'lon': test_globals['location']['longitude'],
'country': test_globals['location']['countryCode'],
'isp': test_globals['ispName'],
},
'ignore_servers': [],
'sizes': {
'upload': [524288, 1048576, 7340032],
},
'counts': {
'upload': 17,
'download': 4
},
'threads': {
'upload': 2,
'download': 8
},
'length': {
'upload': 10,
'download': 10
},
'upload_max': 51,
}
if self._use_socket:
config['sizes']['download'] = [245388, 505544, 1118012, 1986284,
4468241, 7907740, 12407926,
17816816, 24262167, 31625365]
else:
config['sizes']['download'] = [350, 500, 750, 1000, 1500, 2000,
2500, 3000, 3500, 4000]
self.config = config
self.results.client = config['client']
servers = {}
for server in params['serverList']:
if include and int(server.get('id')) not in include:
continue
if exclude and int(server.get('id')) in exclude:
continue
host, port = server['host'].split(':')
server['host'] = (host, int(port))
server['country'] = server['cc']
d = server.pop('distance')
server['d'] = d
try:
servers[d].append(server)
except KeyError:
servers[d] = [server]
self.servers = servers
return self.servers
def get_closest_servers(self, limit=5):
"""Limit servers to the closest speedtest.net servers based on
geographic distance
@ -1532,6 +1626,8 @@ class Speedtest(object):
continue
for _ in range(0, 3):
printer('%s %s' % ('PING', server['host']),
debug=True)
start = timeit.default_timer()
try:
sock.sendall(
@ -1596,6 +1692,10 @@ class Speedtest(object):
for size in self.config['sizes']['download']:
for _ in range(0, self.config['counts']['download']):
requests.append(size)
printer(
'DOWNLOAD %s %s' % (self.best['host'], size),
debug=True
)
request_count = len(requests)
else:
@ -1687,6 +1787,10 @@ class Speedtest(object):
for i, size in enumerate(sizes):
if self._use_socket:
requests.append(size)
printer(
'UPLOAD %s %s' % (self.best['host'], size),
debug=True
)
else:
# We set ``0`` for ``start`` and handle setting the actual
# ``start`` in ``HTTPUploader`` to get better measurements
@ -1847,6 +1951,7 @@ def parse_args():
help='Exclude a server from selection. Can be '
'supplied multiple times')
parser.add_argument('--mini', help='URL of the Speedtest Mini server')
parser.add_argument('--custom', help='URL of the Speedtest Custom Server')
parser.add_argument('--source', help='Source IP address to bind to')
parser.add_argument('--timeout', default=10, type=PARSER_TYPE_FLOAT,
help='HTTP timeout in seconds. Default 10')
@ -1884,6 +1989,7 @@ def validate_optional_args(args):
"""
optional_args = {
'json': ('json/simplejson python module', json),
'custom': ('json/simplejson python module', json),
'secure': ('SSL support', HTTPSConnection),
}
@ -1964,18 +2070,32 @@ def shell():
printer('Retrieving speedtest.net configuration...', quiet)
try:
kwargs = {}
if args.custom:
kwargs['config'] = {}
speedtest = Speedtest(
source_address=args.source,
timeout=args.timeout,
secure=args.secure,
use_socket=args.socket
use_socket=args.socket,
**kwargs
)
except (ConfigRetrievalError,) + HTTP_ERRORS:
printer('Cannot retrieve speedtest configuration', error=True)
raise SpeedtestCLIError(get_exception())
if args.custom:
kwargs = {}
if not args.list:
kwargs.update({
'include': args.server,
'exclude': args.exclude,
})
speedtest.set_custom_server(args.custom, **kwargs)
if args.list:
try:
if not args.custom:
speedtest.get_servers()
except (ServersRetrievalError,) + HTTP_ERRORS:
printer('Cannot retrieve speedtest server list', error=True)
@ -1998,6 +2118,7 @@ def shell():
if not args.mini:
printer('Retrieving speedtest.net server list...', quiet)
if not args.custom:
try:
speedtest.get_servers(servers=args.server, exclude=args.exclude)
except NoMatchedServers: