Advertisement
Guest User

Untitled

a guest
Aug 23rd, 2019
125
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 59.96 KB | None | 0 0
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. # Copyright 2012-2018 Matt Martz
  4. # All Rights Reserved.
  5. #
  6. #    Licensed under the Apache License, Version 2.0 (the "License"); you may
  7. #    not use this file except in compliance with the License. You may obtain
  8. #    a copy of the License at
  9. #
  10. #         http://www.apache.org/licenses/LICENSE-2.0
  11. #
  12. #    Unless required by applicable law or agreed to in writing, software
  13. #    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
  14. #    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
  15. #    License for the specific language governing permissions and limitations
  16. #    under the License.
  17.  
  18. import os
  19. import re
  20. import csv
  21. import sys
  22. import math
  23. import errno
  24. import signal
  25. import socket
  26. import timeit
  27. import datetime
  28. import platform
  29. import threading
  30. import xml.parsers.expat
  31.  
  32. try:
  33.     import gzip
  34.     GZIP_BASE = gzip.GzipFile
  35. except ImportError:
  36.     gzip = None
  37.     GZIP_BASE = object
  38.  
  39. __version__ = '2.0.0'
  40.  
  41.  
  42. class FakeShutdownEvent(object):
  43.     """Class to fake a threading.Event.isSet so that users of this module
  44.    are not required to register their own threading.Event()
  45.    """
  46.  
  47.     @staticmethod
  48.     def isSet():
  49.         "Dummy method to always return false"""
  50.         return False
  51.  
  52.  
  53. # Some global variables we use
  54. DEBUG = False
  55. _GLOBAL_DEFAULT_TIMEOUT = object()
  56.  
  57. # Begin import game to handle Python 2 and Python 3
  58. try:
  59.     import json
  60. except ImportError:
  61.     try:
  62.         import simplejson as json
  63.     except ImportError:
  64.         json = None
  65.  
  66. try:
  67.     import xml.etree.cElementTree as ET
  68. except ImportError:
  69.     try:
  70.         import xml.etree.ElementTree as ET
  71.     except ImportError:
  72.         from xml.dom import minidom as DOM
  73.         ET = None
  74.  
  75. try:
  76.     from urllib2 import (urlopen, Request, HTTPError, URLError,
  77.                          AbstractHTTPHandler, ProxyHandler,
  78.                          HTTPDefaultErrorHandler, HTTPRedirectHandler,
  79.                          HTTPErrorProcessor, OpenerDirector)
  80. except ImportError:
  81.     from urllib.request import (urlopen, Request, HTTPError, URLError,
  82.                                 AbstractHTTPHandler, ProxyHandler,
  83.                                 HTTPDefaultErrorHandler, HTTPRedirectHandler,
  84.                                 HTTPErrorProcessor, OpenerDirector)
  85.  
  86. try:
  87.     from httplib import HTTPConnection
  88. except ImportError:
  89.     from http.client import HTTPConnection
  90.  
  91. try:
  92.     from httplib import HTTPSConnection
  93. except ImportError:
  94.     try:
  95.         from http.client import HTTPSConnection
  96.     except ImportError:
  97.         HTTPSConnection = None
  98.  
  99. try:
  100.     from Queue import Queue
  101. except ImportError:
  102.     from queue import Queue
  103.  
  104. try:
  105.     from urlparse import urlparse
  106. except ImportError:
  107.     from urllib.parse import urlparse
  108.  
  109. try:
  110.     from urlparse import parse_qs
  111. except ImportError:
  112.     try:
  113.         from urllib.parse import parse_qs
  114.     except ImportError:
  115.         from cgi import parse_qs
  116.  
  117. try:
  118.     from hashlib import md5
  119. except ImportError:
  120.     from md5 import md5
  121.  
  122. try:
  123.     from argparse import ArgumentParser as ArgParser
  124.     from argparse import SUPPRESS as ARG_SUPPRESS
  125.     PARSER_TYPE_INT = int
  126.     PARSER_TYPE_STR = str
  127.     PARSER_TYPE_FLOAT = float
  128. except ImportError:
  129.     from optparse import OptionParser as ArgParser
  130.     from optparse import SUPPRESS_HELP as ARG_SUPPRESS
  131.     PARSER_TYPE_INT = 'int'
  132.     PARSER_TYPE_STR = 'string'
  133.     PARSER_TYPE_FLOAT = 'float'
  134.  
  135. try:
  136.     from cStringIO import StringIO
  137.     BytesIO = None
  138. except ImportError:
  139.     try:
  140.         from StringIO import StringIO
  141.         BytesIO = None
  142.     except ImportError:
  143.         from io import StringIO, BytesIO
  144.  
  145. try:
  146.     import __builtin__
  147. except ImportError:
  148.     import builtins
  149.     from io import TextIOWrapper, FileIO
  150.  
  151.     class _Py3Utf8Output(TextIOWrapper):
  152.         """UTF-8 encoded wrapper around stdout for py3, to override
  153.        ASCII stdout
  154.        """
  155.         def __init__(self, f, **kwargs):
  156.             buf = FileIO(f.fileno(), 'w')
  157.             super(_Py3Utf8Output, self).__init__(
  158.                 buf,
  159.                 encoding='utf8',
  160.                 errors='strict'
  161.             )
  162.  
  163.         def write(self, s):
  164.             super(_Py3Utf8Output, self).write(s)
  165.             self.flush()
  166.  
  167.     _py3_print = getattr(builtins, 'print')
  168.     try:
  169.         _py3_utf8_stdout = _Py3Utf8Output(sys.stdout)
  170.         _py3_utf8_stderr = _Py3Utf8Output(sys.stderr)
  171.     except OSError:
  172.         # sys.stdout/sys.stderr is not a compatible stdout/stderr object
  173.         # just use it and hope things go ok
  174.         _py3_utf8_stdout = sys.stdout
  175.         _py3_utf8_stderr = sys.stderr
  176.  
  177.     def to_utf8(v):
  178.         """No-op encode to utf-8 for py3"""
  179.         return v
  180.  
  181.     def print_(*args, **kwargs):
  182.         """Wrapper function for py3 to print, with a utf-8 encoded stdout"""
  183.         if kwargs.get('file') == sys.stderr:
  184.             kwargs['file'] = _py3_utf8_stderr
  185.         else:
  186.             kwargs['file'] = kwargs.get('file', _py3_utf8_stdout)
  187.         _py3_print(*args, **kwargs)
  188. else:
  189.     del __builtin__
  190.  
  191.     def to_utf8(v):
  192.         """Encode value to utf-8 if possible for py2"""
  193.         try:
  194.             return v.encode('utf8', 'strict')
  195.         except AttributeError:
  196.             return v
  197.  
  198.     def print_(*args, **kwargs):
  199.         """The new-style print function for Python 2.4 and 2.5.
  200.  
  201.        Taken from https://pypi.python.org/pypi/six/
  202.  
  203.        Modified to set encoding to UTF-8 always, and to flush after write
  204.        """
  205.         fp = kwargs.pop("file", sys.stdout)
  206.         if fp is None:
  207.             return
  208.  
  209.         def write(data):
  210.             if not isinstance(data, basestring):
  211.                 data = str(data)
  212.             # If the file has an encoding, encode unicode with it.
  213.             encoding = 'utf8'  # Always trust UTF-8 for output
  214.             if (isinstance(fp, file) and
  215.                     isinstance(data, unicode) and
  216.                     encoding is not None):
  217.                 errors = getattr(fp, "errors", None)
  218.                 if errors is None:
  219.                     errors = "strict"
  220.                 data = data.encode(encoding, errors)
  221.             fp.write(data)
  222.             fp.flush()
  223.         want_unicode = False
  224.         sep = kwargs.pop("sep", None)
  225.         if sep is not None:
  226.             if isinstance(sep, unicode):
  227.                 want_unicode = True
  228.             elif not isinstance(sep, str):
  229.                 raise TypeError("sep must be None or a string")
  230.         end = kwargs.pop("end", None)
  231.         if end is not None:
  232.             if isinstance(end, unicode):
  233.                 want_unicode = True
  234.             elif not isinstance(end, str):
  235.                 raise TypeError("end must be None or a string")
  236.         if kwargs:
  237.             raise TypeError("invalid keyword arguments to print()")
  238.         if not want_unicode:
  239.             for arg in args:
  240.                 if isinstance(arg, unicode):
  241.                     want_unicode = True
  242.                     break
  243.         if want_unicode:
  244.             newline = unicode("\n")
  245.             space = unicode(" ")
  246.         else:
  247.             newline = "\n"
  248.             space = " "
  249.         if sep is None:
  250.             sep = space
  251.         if end is None:
  252.             end = newline
  253.         for i, arg in enumerate(args):
  254.             if i:
  255.                 write(sep)
  256.             write(arg)
  257.         write(end)
  258.  
  259.  
  260. # Exception "constants" to support Python 2 through Python 3
  261. try:
  262.     import ssl
  263.     try:
  264.         CERT_ERROR = (ssl.CertificateError,)
  265.     except AttributeError:
  266.         CERT_ERROR = tuple()
  267.  
  268.     HTTP_ERRORS = ((HTTPError, URLError, socket.error, ssl.SSLError) +
  269.                    CERT_ERROR)
  270. except ImportError:
  271.     HTTP_ERRORS = (HTTPError, URLError, socket.error)
  272.  
  273.  
  274. class SpeedtestException(Exception):
  275.     """Base exception for this module"""
  276.  
  277.  
  278. class SpeedtestCLIError(SpeedtestException):
  279.     """Generic exception for raising errors during CLI operation"""
  280.  
  281.  
  282. class SpeedtestHTTPError(SpeedtestException):
  283.     """Base HTTP exception for this module"""
  284.  
  285.  
  286. class SpeedtestConfigError(SpeedtestException):
  287.     """Configuration provided is invalid"""
  288.  
  289.  
  290. class ConfigRetrievalError(SpeedtestHTTPError):
  291.     """Could not retrieve config.php"""
  292.  
  293.  
  294. class ServersRetrievalError(SpeedtestHTTPError):
  295.     """Could not retrieve speedtest-servers.php"""
  296.  
  297.  
  298. class InvalidServerIDType(SpeedtestException):
  299.     """Server ID used for filtering was not an integer"""
  300.  
  301.  
  302. class NoMatchedServers(SpeedtestException):
  303.     """No servers matched when filtering"""
  304.  
  305.  
  306. class SpeedtestMiniConnectFailure(SpeedtestException):
  307.     """Could not connect to the provided speedtest mini server"""
  308.  
  309.  
  310. class InvalidSpeedtestMiniServer(SpeedtestException):
  311.     """Server provided as a speedtest mini server does not actually appear
  312.    to be a speedtest mini server
  313.    """
  314.  
  315.  
  316. class ShareResultsConnectFailure(SpeedtestException):
  317.     """Could not connect to speedtest.net API to POST results"""
  318.  
  319.  
  320. class ShareResultsSubmitFailure(SpeedtestException):
  321.     """Unable to successfully POST results to speedtest.net API after
  322.    connection
  323.    """
  324.  
  325.  
  326. class SpeedtestUploadTimeout(SpeedtestException):
  327.     """testlength configuration reached during upload
  328.    Used to ensure the upload halts when no additional data should be sent
  329.    """
  330.  
  331.  
  332. class SpeedtestBestServerFailure(SpeedtestException):
  333.     """Unable to determine best server"""
  334.  
  335.  
  336. class SpeedtestMissingBestServer(SpeedtestException):
  337.     """get_best_server not called or not able to determine best server"""
  338.  
  339.  
  340. def create_connection(address, timeout=_GLOBAL_DEFAULT_TIMEOUT,
  341.                       source_address=None):
  342.     """Connect to *address* and return the socket object.
  343.  
  344.    Convenience function.  Connect to *address* (a 2-tuple ``(host,
  345.    port)``) and return the socket object.  Passing the optional
  346.    *timeout* parameter will set the timeout on the socket instance
  347.    before attempting to connect.  If no *timeout* is supplied, the
  348.    global default timeout setting returned by :func:`getdefaulttimeout`
  349.    is used.  If *source_address* is set it must be a tuple of (host, port)
  350.    for the socket to bind as a source address before making the connection.
  351.    An host of '' or port 0 tells the OS to use the default.
  352.  
  353.    Largely vendored from Python 2.7, modified to work with Python 2.4
  354.    """
  355.  
  356.     host, port = address
  357.     err = None
  358.     for res in socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM):
  359.         af, socktype, proto, canonname, sa = res
  360.         sock = None
  361.         try:
  362.             sock = socket.socket(af, socktype, proto)
  363.             if timeout is not _GLOBAL_DEFAULT_TIMEOUT:
  364.                 sock.settimeout(float(timeout))
  365.             if source_address:
  366.                 sock.bind(source_address)
  367.             sock.connect(sa)
  368.             return sock
  369.  
  370.         except socket.error:
  371.             err = get_exception()
  372.             if sock is not None:
  373.                 sock.close()
  374.  
  375.     if err is not None:
  376.         raise err
  377.     else:
  378.         raise socket.error("getaddrinfo returns an empty list")
  379.  
  380.  
  381. class SpeedtestHTTPConnection(HTTPConnection):
  382.     """Custom HTTPConnection to support source_address across
  383.    Python 2.4 - Python 3
  384.    """
  385.     def __init__(self, *args, **kwargs):
  386.         source_address = kwargs.pop('source_address', None)
  387.         context = kwargs.pop('context', None)
  388.         timeout = kwargs.pop('timeout', 10)
  389.  
  390.         HTTPConnection.__init__(self, *args, **kwargs)
  391.  
  392.         self.source_address = source_address
  393.         self._context = context
  394.         self.timeout = timeout
  395.  
  396.     def connect(self):
  397.         """Connect to the host and port specified in __init__."""
  398.         try:
  399.             self.sock = socket.create_connection(
  400.                 (self.host, self.port),
  401.                 self.timeout,
  402.                 self.source_address
  403.             )
  404.         except (AttributeError, TypeError):
  405.             self.sock = create_connection(
  406.                 (self.host, self.port),
  407.                 self.timeout,
  408.                 self.source_address
  409.             )
  410.  
  411.  
  412. if HTTPSConnection:
  413.     class SpeedtestHTTPSConnection(HTTPSConnection,
  414.                                    SpeedtestHTTPConnection):
  415.         """Custom HTTPSConnection to support source_address across
  416.        Python 2.4 - Python 3
  417.        """
  418.         def connect(self):
  419.             "Connect to a host on a given (SSL) port."
  420.  
  421.             SpeedtestHTTPConnection.connect(self)
  422.  
  423.             kwargs = {}
  424.             if hasattr(ssl, 'SSLContext'):
  425.                 kwargs['server_hostname'] = self.host
  426.  
  427.             self.sock = self._context.wrap_socket(self.sock, **kwargs)
  428.  
  429.  
  430. def _build_connection(connection, source_address, timeout, context=None):
  431.     """Cross Python 2.4 - Python 3 callable to build an ``HTTPConnection`` or
  432.    ``HTTPSConnection`` with the args we need
  433.  
  434.    Called from ``http(s)_open`` methods of ``SpeedtestHTTPHandler`` or
  435.    ``SpeedtestHTTPSHandler``
  436.    """
  437.     def inner(host, **kwargs):
  438.         kwargs.update({
  439.             'source_address': source_address,
  440.             'timeout': timeout
  441.         })
  442.         if context:
  443.             kwargs['context'] = context
  444.         return connection(host, **kwargs)
  445.     return inner
  446.  
  447.  
  448. class SpeedtestHTTPHandler(AbstractHTTPHandler):
  449.     """Custom ``HTTPHandler`` that can build a ``HTTPConnection`` with the
  450.    args we need for ``source_address`` and ``timeout``
  451.    """
  452.     def __init__(self, debuglevel=0, source_address=None, timeout=10):
  453.         AbstractHTTPHandler.__init__(self, debuglevel)
  454.         self.source_address = source_address
  455.         self.timeout = timeout
  456.  
  457.     def http_open(self, req):
  458.         return self.do_open(
  459.             _build_connection(
  460.                 SpeedtestHTTPConnection,
  461.                 self.source_address,
  462.                 self.timeout
  463.             ),
  464.             req
  465.         )
  466.  
  467.     http_request = AbstractHTTPHandler.do_request_
  468.  
  469.  
  470. class SpeedtestHTTPSHandler(AbstractHTTPHandler):
  471.     """Custom ``HTTPSHandler`` that can build a ``HTTPSConnection`` with the
  472.    args we need for ``source_address`` and ``timeout``
  473.    """
  474.     def __init__(self, debuglevel=0, context=None, source_address=None,
  475.                  timeout=10):
  476.         AbstractHTTPHandler.__init__(self, debuglevel)
  477.         self._context = context
  478.         self.source_address = source_address
  479.         self.timeout = timeout
  480.  
  481.     def https_open(self, req):
  482.         return self.do_open(
  483.             _build_connection(
  484.                 SpeedtestHTTPSConnection,
  485.                 self.source_address,
  486.                 self.timeout,
  487.                 context=self._context,
  488.             ),
  489.             req
  490.         )
  491.  
  492.     https_request = AbstractHTTPHandler.do_request_
  493.  
  494.  
  495. def build_opener(source_address=None, timeout=10):
  496.     """Function similar to ``urllib2.build_opener`` that will build
  497.    an ``OpenerDirector`` with the explicit handlers we want,
  498.    ``source_address`` for binding, ``timeout`` and our custom
  499.    `User-Agent`
  500.    """
  501.  
  502.     printer('Timeout set to %d' % timeout, debug=True)
  503.  
  504.     if source_address:
  505.         source_address_tuple = (source_address, 0)
  506.         printer('Binding to source address: %r' % (source_address_tuple,),
  507.                 debug=True)
  508.     else:
  509.         source_address_tuple = None
  510.  
  511.     handlers = [
  512.         ProxyHandler(),
  513.         SpeedtestHTTPHandler(source_address=source_address_tuple,
  514.                              timeout=timeout),
  515.         SpeedtestHTTPSHandler(source_address=source_address_tuple,
  516.                               timeout=timeout),
  517.         HTTPDefaultErrorHandler(),
  518.         HTTPRedirectHandler(),
  519.         HTTPErrorProcessor()
  520.     ]
  521.  
  522.     opener = OpenerDirector()
  523.     opener.addheaders = [('User-agent', build_user_agent())]
  524.  
  525.     for handler in handlers:
  526.         opener.add_handler(handler)
  527.  
  528.     return opener
  529.  
  530.  
  531. class GzipDecodedResponse(GZIP_BASE):
  532.     """A file-like object to decode a response encoded with the gzip
  533.    method, as described in RFC 1952.
  534.  
  535.    Largely copied from ``xmlrpclib``/``xmlrpc.client`` and modified
  536.    to work for py2.4-py3
  537.    """
  538.     def __init__(self, response):
  539.         # response doesn't support tell() and read(), required by
  540.         # GzipFile
  541.         if not gzip:
  542.             raise SpeedtestHTTPError('HTTP response body is gzip encoded, '
  543.                                      'but gzip support is not available')
  544.         IO = BytesIO or StringIO
  545.         self.io = IO()
  546.         while 1:
  547.             chunk = response.read(1024)
  548.             if len(chunk) == 0:
  549.                 break
  550.             self.io.write(chunk)
  551.         self.io.seek(0)
  552.         gzip.GzipFile.__init__(self, mode='rb', fileobj=self.io)
  553.  
  554.     def close(self):
  555.         try:
  556.             gzip.GzipFile.close(self)
  557.         finally:
  558.             self.io.close()
  559.  
  560.  
  561. def get_exception():
  562.     """Helper function to work with py2.4-py3 for getting the current
  563.    exception in a try/except block
  564.    """
  565.     return sys.exc_info()[1]
  566.  
  567.  
  568. def distance(origin, destination):
  569.     """Determine distance between 2 sets of [lat,lon] in km"""
  570.  
  571.     lat1, lon1 = origin
  572.     lat2, lon2 = destination
  573.     radius = 6371  # km
  574.  
  575.     dlat = math.radians(lat2 - lat1)
  576.     dlon = math.radians(lon2 - lon1)
  577.     a = (math.sin(dlat / 2) * math.sin(dlat / 2) +
  578.          math.cos(math.radians(lat1)) *
  579.          math.cos(math.radians(lat2)) * math.sin(dlon / 2) *
  580.          math.sin(dlon / 2))
  581.     c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
  582.     d = radius * c
  583.  
  584.     return d
  585.  
  586.  
  587. def build_user_agent():
  588.     """Build a Mozilla/5.0 compatible User-Agent string"""
  589.  
  590.     ua_tuple = (
  591.         'Mozilla/5.0',
  592.         '(%s; U; %s; en-us)' % (platform.system(), platform.architecture()[0]),
  593.         'Python/%s' % platform.python_version(),
  594.         '(KHTML, like Gecko)',
  595.         'speedtest-cli/%s' % __version__
  596.     )
  597.     user_agent = ' '.join(ua_tuple)
  598.     printer('User-Agent: %s' % user_agent, debug=True)
  599.     return user_agent
  600.  
  601.  
  602. def build_request(url, data=None, headers=None, bump='0', secure=False):
  603.     """Build a urllib2 request object
  604.  
  605.    This function automatically adds a User-Agent header to all requests
  606.  
  607.    """
  608.  
  609.     if not headers:
  610.         headers = {}
  611.  
  612.     if url[0] == ':':
  613.         scheme = ('http', 'https')[bool(secure)]
  614.         schemed_url = '%s%s' % (scheme, url)
  615.     else:
  616.         schemed_url = url
  617.  
  618.     if '?' in url:
  619.         delim = '&'
  620.     else:
  621.         delim = '?'
  622.  
  623.     # WHO YOU GONNA CALL? CACHE BUSTERS!
  624.     final_url = '%s%sx=%s.%s' % (schemed_url, delim,
  625.                                  int(timeit.time.time() * 1000),
  626.                                  bump)
  627.  
  628.     headers.update({
  629.         'Cache-Control': 'no-cache',
  630.     })
  631.  
  632.     printer('%s %s' % (('GET', 'POST')[bool(data)], final_url),
  633.             debug=True)
  634.  
  635.     return Request(final_url, data=data, headers=headers)
  636.  
  637.  
  638. def catch_request(request, opener=None):
  639.     """Helper function to catch common exceptions encountered when
  640.    establishing a connection with a HTTP/HTTPS request
  641.  
  642.    """
  643.  
  644.     if opener:
  645.         _open = opener.open
  646.     else:
  647.         _open = urlopen
  648.  
  649.     try:
  650.         uh = _open(request)
  651.         return uh, False
  652.     except HTTP_ERRORS:
  653.         e = get_exception()
  654.         return None, e
  655.  
  656.  
  657. def get_response_stream(response):
  658.     """Helper function to return either a Gzip reader if
  659.    ``Content-Encoding`` is ``gzip`` otherwise the response itself
  660.  
  661.    """
  662.  
  663.     try:
  664.         getheader = response.headers.getheader
  665.     except AttributeError:
  666.         getheader = response.getheader
  667.  
  668.     if getheader('content-encoding') == 'gzip':
  669.         return GzipDecodedResponse(response)
  670.  
  671.     return response
  672.  
  673.  
  674. def get_attributes_by_tag_name(dom, tag_name):
  675.     """Retrieve an attribute from an XML document and return it in a
  676.    consistent format
  677.  
  678.    Only used with xml.dom.minidom, which is likely only to be used
  679.    with python versions older than 2.5
  680.    """
  681.     elem = dom.getElementsByTagName(tag_name)[0]
  682.     return dict(list(elem.attributes.items()))
  683.  
  684.  
  685. def print_dots(shutdown_event):
  686.     """Built in callback function used by Thread classes for printing
  687.    status
  688.    """
  689.     def inner(current, total, start=False, end=False):
  690.         if shutdown_event.isSet():
  691.             return
  692.  
  693.         sys.stdout.write('.')
  694.         if current + 1 == total and end is True:
  695.             sys.stdout.write('\n')
  696.         sys.stdout.flush()
  697.     return inner
  698.  
  699.  
  700. def do_nothing(*args, **kwargs):
  701.     pass
  702.  
  703.  
  704. class HTTPDownloader(threading.Thread):
  705.     """Thread class for retrieving a URL"""
  706.  
  707.     def __init__(self, i, request, start, timeout, opener=None,
  708.                  shutdown_event=None):
  709.         threading.Thread.__init__(self)
  710.         self.request = request
  711.         self.result = [0]
  712.         self.starttime = start
  713.         self.timeout = timeout
  714.         self.i = i
  715.         if opener:
  716.             self._opener = opener.open
  717.         else:
  718.             self._opener = urlopen
  719.  
  720.         if shutdown_event:
  721.             self._shutdown_event = shutdown_event
  722.         else:
  723.             self._shutdown_event = FakeShutdownEvent()
  724.  
  725.     def run(self):
  726.         try:
  727.             if (timeit.default_timer() - self.starttime) <= self.timeout:
  728.                 f = self._opener(self.request)
  729.                 while (not self._shutdown_event.isSet() and
  730.                         (timeit.default_timer() - self.starttime) <=
  731.                         self.timeout):
  732.                     self.result.append(len(f.read(10240)))
  733.                     if self.result[-1] == 0:
  734.                         break
  735.                 f.close()
  736.         except IOError:
  737.             pass
  738.  
  739.  
  740. class HTTPUploaderData(object):
  741.     """File like object to improve cutting off the upload once the timeout
  742.    has been reached
  743.    """
  744.  
  745.     def __init__(self, length, start, timeout, shutdown_event=None):
  746.         self.length = length
  747.         self.start = start
  748.         self.timeout = timeout
  749.  
  750.         if shutdown_event:
  751.             self._shutdown_event = shutdown_event
  752.         else:
  753.             self._shutdown_event = FakeShutdownEvent()
  754.  
  755.         self._data = None
  756.  
  757.         self.total = [0]
  758.  
  759.     def pre_allocate(self):
  760.         chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
  761.         multiplier = int(round(int(self.length) / 36.0))
  762.         IO = BytesIO or StringIO
  763.         try:
  764.             self._data = IO(
  765.                 ('content1=%s' %
  766.                  (chars * multiplier)[0:int(self.length) - 9]
  767.                  ).encode()
  768.             )
  769.         except MemoryError:
  770.             raise SpeedtestCLIError(
  771.                 'Insufficient memory to pre-allocate upload data. Please '
  772.                 'use --no-pre-allocate'
  773.             )
  774.  
  775.     @property
  776.     def data(self):
  777.         if not self._data:
  778.             self.pre_allocate()
  779.         return self._data
  780.  
  781.     def read(self, n=10240):
  782.         if ((timeit.default_timer() - self.start) <= self.timeout and
  783.                 not self._shutdown_event.isSet()):
  784.             chunk = self.data.read(n)
  785.             self.total.append(len(chunk))
  786.             return chunk
  787.         else:
  788.             raise SpeedtestUploadTimeout()
  789.  
  790.     def __len__(self):
  791.         return self.length
  792.  
  793.  
  794. class HTTPUploader(threading.Thread):
  795.     """Thread class for putting a URL"""
  796.  
  797.     def __init__(self, i, request, start, size, timeout, opener=None,
  798.                  shutdown_event=None):
  799.         threading.Thread.__init__(self)
  800.         self.request = request
  801.         self.request.data.start = self.starttime = start
  802.         self.size = size
  803.         self.result = None
  804.         self.timeout = timeout
  805.         self.i = i
  806.  
  807.         if opener:
  808.             self._opener = opener.open
  809.         else:
  810.             self._opener = urlopen
  811.  
  812.         if shutdown_event:
  813.             self._shutdown_event = shutdown_event
  814.         else:
  815.             self._shutdown_event = FakeShutdownEvent()
  816.  
  817.     def run(self):
  818.         request = self.request
  819.         try:
  820.             if ((timeit.default_timer() - self.starttime) <= self.timeout and
  821.                     not self._shutdown_event.isSet()):
  822.                 try:
  823.                     f = self._opener(request)
  824.                 except TypeError:
  825.                     # PY24 expects a string or buffer
  826.                     # This also causes issues with Ctrl-C, but we will concede
  827.                     # for the moment that Ctrl-C on PY24 isn't immediate
  828.                     request = build_request(self.request.get_full_url(),
  829.                                             data=request.data.read(self.size))
  830.                     f = self._opener(request)
  831.                 f.read(11)
  832.                 f.close()
  833.                 self.result = sum(self.request.data.total)
  834.             else:
  835.                 self.result = 0
  836.         except (IOError, SpeedtestUploadTimeout):
  837.             self.result = sum(self.request.data.total)
  838.  
  839.  
  840. class SpeedtestResults(object):
  841.     """Class for holding the results of a speedtest, including:
  842.  
  843.    Download speed
  844.    Upload speed
  845.    Ping/Latency to test server
  846.    Data about server that the test was run against
  847.  
  848.    Additionally this class can return a result data as a dictionary or CSV,
  849.    as well as submit a POST of the result data to the speedtest.net API
  850.    to get a share results image link.
  851.    """
  852.  
  853.     def __init__(self, download=0, upload=0, ping=0, server=None, client=None,
  854.                  opener=None, secure=False):
  855.         self.download = download
  856.         self.upload = upload
  857.         self.ping = ping
  858.         if server is None:
  859.             self.server = {}
  860.         else:
  861.             self.server = server
  862.         self.client = client or {}
  863.  
  864.         self._share = None
  865.         self.timestamp = '%sZ' % datetime.datetime.utcnow().isoformat()
  866.         self.bytes_received = 0
  867.         self.bytes_sent = 0
  868.  
  869.         if opener:
  870.             self._opener = opener
  871.         else:
  872.             self._opener = build_opener()
  873.  
  874.         self._secure = secure
  875.  
  876.     def __repr__(self):
  877.         return repr(self.dict())
  878.  
  879.     def share(self):
  880.         """POST data to the speedtest.net API to obtain a share results
  881.        link
  882.        """
  883.  
  884.         if self._share:
  885.             return self._share
  886.  
  887.         download = int(round(self.download / 1000.0, 0))
  888.         ping = int(round(self.ping, 0))
  889.         upload = int(round(self.upload / 1000.0, 0))
  890.  
  891.         # Build the request to send results back to speedtest.net
  892.         # We use a list instead of a dict because the API expects parameters
  893.         # in a certain order
  894.         api_data = [
  895.             'recommendedserverid=%s' % self.server['id'],
  896.             'ping=%s' % ping,
  897.             'screenresolution=',
  898.             'promo=',
  899.             'download=%s' % download,
  900.             'screendpi=',
  901.             'upload=%s' % upload,
  902.             'testmethod=http',
  903.             'hash=%s' % md5(('%s-%s-%s-%s' %
  904.                              (ping, upload, download, '297aae72'))
  905.                             .encode()).hexdigest(),
  906.             'touchscreen=none',
  907.             'startmode=pingselect',
  908.             'accuracy=1',
  909.             'bytesreceived=%s' % self.bytes_received,
  910.             'bytessent=%s' % self.bytes_sent,
  911.             'serverid=%s' % self.server['id'],
  912.         ]
  913.  
  914.         headers = {'Referer': 'http://c.speedtest.net/flash/speedtest.swf'}
  915.         request = build_request('://www.speedtest.net/api/api.php',
  916.                                 data='&'.join(api_data).encode(),
  917.                                 headers=headers, secure=self._secure)
  918.         f, e = catch_request(request, opener=self._opener)
  919.         if e:
  920.             raise ShareResultsConnectFailure(e)
  921.  
  922.         response = f.read()
  923.         code = f.code
  924.         f.close()
  925.  
  926.         if int(code) != 200:
  927.             raise ShareResultsSubmitFailure('Could not submit results to '
  928.                                             'speedtest.net')
  929.  
  930.         qsargs = parse_qs(response.decode())
  931.         resultid = qsargs.get('resultid')
  932.         if not resultid or len(resultid) != 1:
  933.             raise ShareResultsSubmitFailure('Could not submit results to '
  934.                                             'speedtest.net')
  935.  
  936.         self._share = 'http://www.speedtest.net/result/%s.png' % resultid[0]
  937.  
  938.         return self._share
  939.  
  940.     def dict(self):
  941.         """Return dictionary of result data"""
  942.  
  943.         return {
  944.             'download': self.download,
  945.             'upload': self.upload,
  946.             'ping': self.ping,
  947.             'server': self.server,
  948.             'timestamp': self.timestamp,
  949.             'bytes_sent': self.bytes_sent,
  950.             'bytes_received': self.bytes_received,
  951.             'share': self._share,
  952.             'client': self.client,
  953.         }
  954.  
  955.     @staticmethod
  956.     def csv_header(delimiter=','):
  957.         """Return CSV Headers"""
  958.  
  959.         row = ['Server ID', 'Sponsor', 'Server Name', 'Timestamp', 'Distance',
  960.                'Ping', 'Download', 'Upload', 'Share', 'IP Address']
  961.         out = StringIO()
  962.         writer = csv.writer(out, delimiter=delimiter, lineterminator='')
  963.         writer.writerow([to_utf8(v) for v in row])
  964.         return out.getvalue()
  965.  
  966.     def csv(self, delimiter=','):
  967.         """Return data in CSV format"""
  968.  
  969.         data = self.dict()
  970.         out = StringIO()
  971.         writer = csv.writer(out, delimiter=delimiter, lineterminator='')
  972.         row = [data['server']['id'], data['server']['sponsor'],
  973.                data['server']['name'], data['timestamp'],
  974.                data['server']['d'], data['ping'], data['download'],
  975.                data['upload'], self._share or '', self.client['ip']]
  976.         writer.writerow([to_utf8(v) for v in row])
  977.         return out.getvalue()
  978.  
  979.     def json(self, pretty=False):
  980.         """Return data in JSON format"""
  981.  
  982.         kwargs = {}
  983.         if pretty:
  984.             kwargs.update({
  985.                 'indent': 4,
  986.                 'sort_keys': True
  987.             })
  988.         return json.dumps(self.dict(), **kwargs)
  989.  
  990.  
  991. class Speedtest(object):
  992.     """Class for performing standard speedtest.net testing operations"""
  993.  
  994.     def __init__(self, config=None, source_address=None, timeout=10,
  995.                  secure=False, shutdown_event=None):
  996.         self.config = {}
  997.  
  998.         self._source_address = source_address
  999.         self._timeout = timeout
  1000.         self._opener = build_opener(source_address, timeout)
  1001.  
  1002.         self._secure = secure
  1003.  
  1004.         if shutdown_event:
  1005.             self._shutdown_event = shutdown_event
  1006.         else:
  1007.             self._shutdown_event = FakeShutdownEvent()
  1008.  
  1009.         self.get_config()
  1010.         if config is not None:
  1011.             self.config.update(config)
  1012.  
  1013.         self.servers = {}
  1014.         self.closest = []
  1015.         self._best = {}
  1016.  
  1017.         self.results = SpeedtestResults(
  1018.             client=self.config['client'],
  1019.             opener=self._opener,
  1020.             secure=secure,
  1021.         )
  1022.  
  1023.     @property
  1024.     def best(self):
  1025.         if not self._best:
  1026.             raise SpeedtestMissingBestServer(
  1027.                 'get_best_server not called or not able to determine best '
  1028.                 'server'
  1029.             )
  1030.         return self._best
  1031.  
  1032.     def get_config(self):
  1033.         """Download the speedtest.net configuration and return only the data
  1034.        we are interested in
  1035.        """
  1036.  
  1037.         headers = {}
  1038.         if gzip:
  1039.             headers['Accept-Encoding'] = 'gzip'
  1040.         request = build_request('://www.speedtest.net/speedtest-config.php',
  1041.                                 headers=headers, secure=self._secure)
  1042.         uh, e = catch_request(request, opener=self._opener)
  1043.         if e:
  1044.             raise ConfigRetrievalError(e)
  1045.         configxml = []
  1046.  
  1047.         stream = get_response_stream(uh)
  1048.  
  1049.         while 1:
  1050.             try:
  1051.                 configxml.append(stream.read(1024))
  1052.             except (OSError, EOFError):
  1053.                 raise ConfigRetrievalError(get_exception())
  1054.             if len(configxml[-1]) == 0:
  1055.                 break
  1056.         stream.close()
  1057.         uh.close()
  1058.  
  1059.         if int(uh.code) != 200:
  1060.             return None
  1061.  
  1062.         printer('Config XML:\n%s' % ''.encode().join(configxml), debug=True)
  1063.  
  1064.         try:
  1065.             root = ET.fromstring(''.encode().join(configxml))
  1066.             server_config = root.find('server-config').attrib
  1067.             download = root.find('download').attrib
  1068.             upload = root.find('upload').attrib
  1069.             # times = root.find('times').attrib
  1070.             client = root.find('client').attrib
  1071.  
  1072.         except AttributeError:
  1073.             root = DOM.parseString(''.join(configxml))
  1074.             server_config = get_attributes_by_tag_name(root, 'server-config')
  1075.             download = get_attributes_by_tag_name(root, 'download')
  1076.             upload = get_attributes_by_tag_name(root, 'upload')
  1077.             # times = get_attributes_by_tag_name(root, 'times')
  1078.             client = get_attributes_by_tag_name(root, 'client')
  1079.  
  1080.         ignore_servers = list(
  1081.             map(int, server_config['ignoreids'].split(','))
  1082.         )
  1083.  
  1084.         ratio = int(upload['ratio'])
  1085.         upload_max = int(upload['maxchunkcount'])
  1086.         up_sizes = [32768, 65536, 131072, 262144, 524288, 1048576, 7340032]
  1087.         sizes = {
  1088.             'upload': up_sizes[ratio - 1:],
  1089.             'download': [350, 500, 750, 1000, 1500, 2000, 2500,
  1090.                          3000, 3500, 4000]
  1091.         }
  1092.  
  1093.         size_count = len(sizes['upload'])
  1094.  
  1095.         upload_count = int(math.ceil(upload_max / size_count))
  1096.  
  1097.         counts = {
  1098.             'upload': upload_count,
  1099.             'download': int(download['threadsperurl'])
  1100.         }
  1101.  
  1102.         threads = {
  1103.             'upload': int(upload['threads']),
  1104.             'download': int(server_config['threadcount']) * 2
  1105.         }
  1106.  
  1107.         length = {
  1108.             'upload': int(upload['testlength']),
  1109.             'download': int(download['testlength'])
  1110.         }
  1111.  
  1112.         self.config.update({
  1113.             'client': client,
  1114.             'ignore_servers': ignore_servers,
  1115.             'sizes': sizes,
  1116.             'counts': counts,
  1117.             'threads': threads,
  1118.             'length': length,
  1119.             'upload_max': upload_count * size_count
  1120.         })
  1121.  
  1122.         self.lat_lon = (float(client['lat']), float(client['lon']))
  1123.  
  1124.         printer('Config:\n%r' % self.config, debug=True)
  1125.  
  1126.         return self.config
  1127.  
  1128.     def get_servers(self, servers=None, exclude=None):
  1129.         """Retrieve a the list of speedtest.net servers, optionally filtered
  1130.        to servers matching those specified in the ``servers`` argument
  1131.        """
  1132.         if servers is None:
  1133.             servers = []
  1134.  
  1135.         if exclude is None:
  1136.             exclude = []
  1137.  
  1138.         self.servers.clear()
  1139.  
  1140.         for server_list in (servers, exclude):
  1141.             for i, s in enumerate(server_list):
  1142.                 try:
  1143.                     server_list[i] = int(s)
  1144.                 except ValueError:
  1145.                     raise InvalidServerIDType(
  1146.                         '%s is an invalid server type, must be int' % s
  1147.                     )
  1148.  
  1149.         urls = [
  1150.             '://www.speedtest.net/speedtest-servers-static.php',
  1151.             'http://c.speedtest.net/speedtest-servers-static.php',
  1152.             '://www.speedtest.net/speedtest-servers.php',
  1153.             'http://c.speedtest.net/speedtest-servers.php',
  1154.         ]
  1155.  
  1156.         headers = {}
  1157.         if gzip:
  1158.             headers['Accept-Encoding'] = 'gzip'
  1159.  
  1160.         errors = []
  1161.         for url in urls:
  1162.             try:
  1163.                 request = build_request(
  1164.                     '%s?threads=%s' % (url,
  1165.                                        self.config['threads']['download']),
  1166.                     headers=headers,
  1167.                     secure=self._secure
  1168.                 )
  1169.                 uh, e = catch_request(request, opener=self._opener)
  1170.                 if e:
  1171.                     errors.append('%s' % e)
  1172.                     raise ServersRetrievalError()
  1173.  
  1174.                 stream = get_response_stream(uh)
  1175.  
  1176.                 serversxml = []
  1177.                 while 1:
  1178.                     try:
  1179.                         serversxml.append(stream.read(1024))
  1180.                     except (OSError, EOFError):
  1181.                         raise ServersRetrievalError(get_exception())
  1182.                     if len(serversxml[-1]) == 0:
  1183.                         break
  1184.  
  1185.                 stream.close()
  1186.                 uh.close()
  1187.  
  1188.                 if int(uh.code) != 200:
  1189.                     raise ServersRetrievalError()
  1190.  
  1191.                 printer('Servers XML:\n%s' % ''.encode().join(serversxml),
  1192.                         debug=True)
  1193.  
  1194.                 try:
  1195.                     try:
  1196.                         root = ET.fromstring(''.encode().join(serversxml))
  1197.                         elements = root.getiterator('server')
  1198.                     except AttributeError:
  1199.                         root = DOM.parseString(''.join(serversxml))
  1200.                         elements = root.getElementsByTagName('server')
  1201.                 except (SyntaxError, xml.parsers.expat.ExpatError):
  1202.                     raise ServersRetrievalError()
  1203.  
  1204.                 for server in elements:
  1205.                     try:
  1206.                         attrib = server.attrib
  1207.                     except AttributeError:
  1208.                         attrib = dict(list(server.attributes.items()))
  1209.  
  1210.                     if servers and int(attrib.get('id')) not in servers:
  1211.                         continue
  1212.  
  1213.                     if (int(attrib.get('id')) in self.config['ignore_servers']
  1214.                             or int(attrib.get('id')) in exclude):
  1215.                         continue
  1216.  
  1217.                     try:
  1218.                         d = distance(self.lat_lon,
  1219.                                      (float(attrib.get('lat')),
  1220.                                       float(attrib.get('lon'))))
  1221.                     except Exception:
  1222.                         continue
  1223.  
  1224.                     attrib['d'] = d
  1225.  
  1226.                     try:
  1227.                         self.servers[d].append(attrib)
  1228.                     except KeyError:
  1229.                         self.servers[d] = [attrib]
  1230.  
  1231.                 break
  1232.  
  1233.             except ServersRetrievalError:
  1234.                 continue
  1235.  
  1236.         if (servers or exclude) and not self.servers:
  1237.             raise NoMatchedServers()
  1238.  
  1239.         return self.servers
  1240.  
  1241.     def set_mini_server(self, server):
  1242.         """Instead of querying for a list of servers, set a link to a
  1243.        speedtest mini server
  1244.        """
  1245.  
  1246.         urlparts = urlparse(server)
  1247.  
  1248.         name, ext = os.path.splitext(urlparts[2])
  1249.         if ext:
  1250.             url = os.path.dirname(server)
  1251.         else:
  1252.             url = server
  1253.  
  1254.         request = build_request(url)
  1255.         uh, e = catch_request(request, opener=self._opener)
  1256.         if e:
  1257.             raise SpeedtestMiniConnectFailure('Failed to connect to %s' %
  1258.                                               server)
  1259.         else:
  1260.             text = uh.read()
  1261.             uh.close()
  1262.  
  1263.         extension = re.findall('upload_?[Ee]xtension: "([^"]+)"',
  1264.                                text.decode())
  1265.         if not extension:
  1266.             for ext in ['php', 'asp', 'aspx', 'jsp']:
  1267.                 try:
  1268.                     f = self._opener.open(
  1269.                         '%s/speedtest/upload.%s' % (url, ext)
  1270.                     )
  1271.                 except Exception:
  1272.                     pass
  1273.                 else:
  1274.                     data = f.read().strip().decode()
  1275.                     if (f.code == 200 and
  1276.                             len(data.splitlines()) == 1 and
  1277.                             re.match('size=[0-9]', data)):
  1278.                         extension = [ext]
  1279.                         break
  1280.         if not urlparts or not extension:
  1281.             raise InvalidSpeedtestMiniServer('Invalid Speedtest Mini Server: '
  1282.                                              '%s' % server)
  1283.  
  1284.         self.servers = [{
  1285.             'sponsor': 'Speedtest Mini',
  1286.             'name': urlparts[1],
  1287.             'd': 0,
  1288.             'url': '%s/speedtest/upload.%s' % (url.rstrip('/'), extension[0]),
  1289.             'latency': 0,
  1290.             'id': 0
  1291.         }]
  1292.  
  1293.         return self.servers
  1294.  
  1295.     def get_closest_servers(self, limit=5):
  1296.         """Limit servers to the closest speedtest.net servers based on
  1297.        geographic distance
  1298.        """
  1299.  
  1300.         if not self.servers:
  1301.             self.get_servers()
  1302.  
  1303.         for d in sorted(self.servers.keys()):
  1304.             for s in self.servers[d]:
  1305.                 self.closest.append(s)
  1306.                 if len(self.closest) == limit:
  1307.                     break
  1308.             else:
  1309.                 continue
  1310.             break
  1311.  
  1312.         printer('Closest Servers:\n%r' % self.closest, debug=True)
  1313.         return self.closest
  1314.  
  1315.     def get_best_server(self, servers=None):
  1316.         """Perform a speedtest.net "ping" to determine which speedtest.net
  1317.        server has the lowest latency
  1318.        """
  1319.  
  1320.         if not servers:
  1321.             if not self.closest:
  1322.                 servers = self.get_closest_servers()
  1323.             servers = self.closest
  1324.  
  1325.         if self._source_address:
  1326.             source_address_tuple = (self._source_address, 0)
  1327.         else:
  1328.             source_address_tuple = None
  1329.  
  1330.         user_agent = build_user_agent()
  1331.  
  1332.         results = {}
  1333.         for server in servers:
  1334.             cum = []
  1335.             url = os.path.dirname(server['url'])
  1336.             stamp = int(timeit.time.time() * 1000)
  1337.             latency_url = '%s/latency.txt?x=%s' % (url, stamp)
  1338.             for i in range(0, 3):
  1339.                 this_latency_url = '%s.%s' % (latency_url, i)
  1340.                 printer('%s %s' % ('GET', this_latency_url),
  1341.                         debug=True)
  1342.                 urlparts = urlparse(latency_url)
  1343.                 try:
  1344.                     if urlparts[0] == 'https':
  1345.                         h = SpeedtestHTTPSConnection(
  1346.                             urlparts[1],
  1347.                             source_address=source_address_tuple
  1348.                         )
  1349.                     else:
  1350.                         h = SpeedtestHTTPConnection(
  1351.                             urlparts[1],
  1352.                             source_address=source_address_tuple
  1353.                         )
  1354.                     headers = {'User-Agent': user_agent}
  1355.                     path = '%s?%s' % (urlparts[2], urlparts[4])
  1356.                     start = timeit.default_timer()
  1357.                     h.request("GET", path, headers=headers)
  1358.                     r = h.getresponse()
  1359.                     total = (timeit.default_timer() - start)
  1360.                 except HTTP_ERRORS:
  1361.                     e = get_exception()
  1362.                     printer('ERROR: %r' % e, debug=True)
  1363.                     cum.append(3600)
  1364.                     continue
  1365.  
  1366.                 text = r.read(9)
  1367.                 if int(r.status) == 200 and text == 'test=test'.encode():
  1368.                     cum.append(total)
  1369.                 else:
  1370.                     cum.append(3600)
  1371.                 h.close()
  1372.  
  1373.             avg = round((sum(cum) / 6) * 1000.0, 3)
  1374.             results[avg] = server
  1375.  
  1376.         try:
  1377.             fastest = sorted(results.keys())[0]
  1378.         except IndexError:
  1379.             raise SpeedtestBestServerFailure('Unable to connect to servers to '
  1380.                                              'test latency.')
  1381.         best = results[fastest]
  1382.         best['latency'] = fastest
  1383.  
  1384.         self.results.ping = fastest
  1385.         self.results.server = best
  1386.  
  1387.         self._best.update(best)
  1388.         printer('Best Server:\n%r' % best, debug=True)
  1389.         return best
  1390.  
  1391.     def download(self, callback=do_nothing):
  1392.         """Test download speed against speedtest.net"""
  1393.  
  1394.         urls = []
  1395.         for size in self.config['sizes']['download']:
  1396.             for _ in range(0, self.config['counts']['download']):
  1397.                 urls.append('%s/random%sx%s.jpg' %
  1398.                             (os.path.dirname(self.best['url']), size, size))
  1399.  
  1400.         request_count = len(urls)
  1401.         requests = []
  1402.         for i, url in enumerate(urls):
  1403.             requests.append(
  1404.                 build_request(url, bump=i, secure=self._secure)
  1405.             )
  1406.  
  1407.         def producer(q, requests, request_count):
  1408.             for i, request in enumerate(requests):
  1409.                 thread = HTTPDownloader(
  1410.                     i,
  1411.                     request,
  1412.                     start,
  1413.                     self.config['length']['download'],
  1414.                     opener=self._opener,
  1415.                     shutdown_event=self._shutdown_event
  1416.                 )
  1417.                 thread.start()
  1418.                 q.put(thread, True)
  1419.                 callback(i, request_count, start=True)
  1420.  
  1421.         finished = []
  1422.  
  1423.         def consumer(q, request_count):
  1424.             while len(finished) < request_count:
  1425.                 thread = q.get(True)
  1426.                 while thread.isAlive():
  1427.                     thread.join(timeout=0.1)
  1428.                 finished.append(sum(thread.result))
  1429.                 callback(thread.i, request_count, end=True)
  1430.  
  1431.         q = Queue(self.config['threads']['download'])
  1432.         prod_thread = threading.Thread(target=producer,
  1433.                                        args=(q, requests, request_count))
  1434.         cons_thread = threading.Thread(target=consumer,
  1435.                                        args=(q, request_count))
  1436.         start = timeit.default_timer()
  1437.         prod_thread.start()
  1438.         cons_thread.start()
  1439.         while prod_thread.isAlive():
  1440.             prod_thread.join(timeout=0.1)
  1441.         while cons_thread.isAlive():
  1442.             cons_thread.join(timeout=0.1)
  1443.  
  1444.         stop = timeit.default_timer()
  1445.         self.results.bytes_received = sum(finished)
  1446.         self.results.download = (
  1447.             (self.results.bytes_received / (stop - start)) * 8.0
  1448.         )
  1449.         if self.results.download > 100000:
  1450.             self.config['threads']['upload'] = 8
  1451.         return self.results.download
  1452.  
  1453.     def upload(self, callback=do_nothing, pre_allocate=True):
  1454.         """Test upload speed against speedtest.net"""
  1455.  
  1456.         sizes = []
  1457.  
  1458.         for size in self.config['sizes']['upload']:
  1459.             for _ in range(0, self.config['counts']['upload']):
  1460.                 sizes.append(size)
  1461.  
  1462.         # request_count = len(sizes)
  1463.         request_count = self.config['upload_max']
  1464.  
  1465.         requests = []
  1466.         for i, size in enumerate(sizes):
  1467.             # We set ``0`` for ``start`` and handle setting the actual
  1468.             # ``start`` in ``HTTPUploader`` to get better measurements
  1469.             data = HTTPUploaderData(
  1470.                 size,
  1471.                 0,
  1472.                 self.config['length']['upload'],
  1473.                 shutdown_event=self._shutdown_event
  1474.             )
  1475.             if pre_allocate:
  1476.                 data.pre_allocate()
  1477.             requests.append(
  1478.                 (
  1479.                     build_request(self.best['url'], data, secure=self._secure),
  1480.                     size
  1481.                 )
  1482.             )
  1483.  
  1484.         def producer(q, requests, request_count):
  1485.             for i, request in enumerate(requests[:request_count]):
  1486.                 thread = HTTPUploader(
  1487.                     i,
  1488.                     request[0],
  1489.                     start,
  1490.                     request[1],
  1491.                     self.config['length']['upload'],
  1492.                     opener=self._opener,
  1493.                     shutdown_event=self._shutdown_event
  1494.                 )
  1495.                 thread.start()
  1496.                 q.put(thread, True)
  1497.                 callback(i, request_count, start=True)
  1498.  
  1499.         finished = []
  1500.  
  1501.         def consumer(q, request_count):
  1502.             while len(finished) < request_count:
  1503.                 thread = q.get(True)
  1504.                 while thread.isAlive():
  1505.                     thread.join(timeout=0.1)
  1506.                 finished.append(thread.result)
  1507.                 callback(thread.i, request_count, end=True)
  1508.  
  1509.         q = Queue(self.config['threads']['upload'])
  1510.         prod_thread = threading.Thread(target=producer,
  1511.                                        args=(q, requests, request_count))
  1512.         cons_thread = threading.Thread(target=consumer,
  1513.                                        args=(q, request_count))
  1514.         start = timeit.default_timer()
  1515.         prod_thread.start()
  1516.         cons_thread.start()
  1517.         while prod_thread.isAlive():
  1518.             prod_thread.join(timeout=0.1)
  1519.         while cons_thread.isAlive():
  1520.             cons_thread.join(timeout=0.1)
  1521.  
  1522.         stop = timeit.default_timer()
  1523.         self.results.bytes_sent = sum(finished)
  1524.         self.results.upload = (
  1525.             (self.results.bytes_sent / (stop - start)) * 8.0
  1526.         )
  1527.         return self.results.upload
  1528.  
  1529.  
  1530. def ctrl_c(shutdown_event):
  1531.     """Catch Ctrl-C key sequence and set a SHUTDOWN_EVENT for our threaded
  1532.    operations
  1533.    """
  1534.     def inner(signum, frame):
  1535.         shutdown_event.set()
  1536.         printer('\nCancelling...', error=True)
  1537.         sys.exit(0)
  1538.     return inner
  1539.  
  1540.  
  1541. def version():
  1542.     """Print the version"""
  1543.  
  1544.     printer(__version__)
  1545.     sys.exit(0)
  1546.  
  1547.  
  1548. def csv_header(delimiter=','):
  1549.     """Print the CSV Headers"""
  1550.  
  1551.     printer(SpeedtestResults.csv_header(delimiter=delimiter))
  1552.     sys.exit(0)
  1553.  
  1554.  
  1555. def parse_args():
  1556.     """Function to handle building and parsing of command line arguments"""
  1557.     description = (
  1558.         'Command line interface for testing internet bandwidth using '
  1559.         'speedtest.net.\n'
  1560.         '------------------------------------------------------------'
  1561.         '--------------\n'
  1562.         'https://github.com/sivel/speedtest-cli')
  1563.  
  1564.     parser = ArgParser(description=description)
  1565.     # Give optparse.OptionParser an `add_argument` method for
  1566.     # compatibility with argparse.ArgumentParser
  1567.     try:
  1568.         parser.add_argument = parser.add_option
  1569.     except AttributeError:
  1570.         pass
  1571.     parser.add_argument('--no-download', dest='download', default=True,
  1572.                         action='store_const', const=False,
  1573.                         help='Do not perform download test')
  1574.     parser.add_argument('--no-upload', dest='upload', default=True,
  1575.                         action='store_const', const=False,
  1576.                         help='Do not perform upload test')
  1577.     parser.add_argument('--bytes', dest='units', action='store_const',
  1578.                         const=('byte', 8), default=('bit', 1),
  1579.                         help='Display values in bytes instead of bits. Does '
  1580.                              'not affect the image generated by --share, nor '
  1581.                              'output from --json or --csv')
  1582.     parser.add_argument('--share', action='store_true',
  1583.                         help='Generate and provide a URL to the speedtest.net '
  1584.                              'share results image, not displayed with --csv')
  1585.     parser.add_argument('--simple', action='store_true', default=False,
  1586.                         help='Suppress verbose output, only show basic '
  1587.                              'information')
  1588.     parser.add_argument('--csv', action='store_true', default=False,
  1589.                         help='Suppress verbose output, only show basic '
  1590.                              'information in CSV format. Speeds listed in '
  1591.                              'bit/s and not affected by --bytes')
  1592.     parser.add_argument('--csv-delimiter', default=',', type=PARSER_TYPE_STR,
  1593.                         help='Single character delimiter to use in CSV '
  1594.                              'output. Default ","')
  1595.     parser.add_argument('--csv-header', action='store_true', default=False,
  1596.                         help='Print CSV headers')
  1597.     parser.add_argument('--json', action='store_true', default=False,
  1598.                         help='Suppress verbose output, only show basic '
  1599.                              'information in JSON format. Speeds listed in '
  1600.                              'bit/s and not affected by --bytes')
  1601.     parser.add_argument('--list', action='store_true',
  1602.                         help='Display a list of speedtest.net servers '
  1603.                              'sorted by distance')
  1604.     parser.add_argument('--server', type=PARSER_TYPE_INT, action='append',
  1605.                         help='Specify a server ID to test against. Can be '
  1606.                              'supplied multiple times')
  1607.     parser.add_argument('--exclude', type=PARSER_TYPE_INT, action='append',
  1608.                         help='Exclude a server from selection. Can be '
  1609.                              'supplied multiple times')
  1610.     parser.add_argument('--mini', help='URL of the Speedtest Mini server')
  1611.     parser.add_argument('--source', help='Source IP address to bind to')
  1612.     parser.add_argument('--timeout', default=10, type=PARSER_TYPE_FLOAT,
  1613.                         help='HTTP timeout in seconds. Default 10')
  1614.     parser.add_argument('--secure', action='store_true',
  1615.                         help='Use HTTPS instead of HTTP when communicating '
  1616.                              'with speedtest.net operated servers')
  1617.     parser.add_argument('--no-pre-allocate', dest='pre_allocate',
  1618.                         action='store_const', default=True, const=False,
  1619.                         help='Do not pre allocate upload data. Pre allocation '
  1620.                              'is enabled by default to improve upload '
  1621.                              'performance. To support systems with '
  1622.                              'insufficient memory, use this option to avoid a '
  1623.                              'MemoryError')
  1624.     parser.add_argument('--version', action='store_true',
  1625.                         help='Show the version number and exit')
  1626.     parser.add_argument('--debug', action='store_true',
  1627.                         help=ARG_SUPPRESS, default=ARG_SUPPRESS)
  1628.  
  1629.     options = parser.parse_args()
  1630.     if isinstance(options, tuple):
  1631.         args = options[0]
  1632.     else:
  1633.         args = options
  1634.     return args
  1635.  
  1636.  
  1637. def validate_optional_args(args):
  1638.     """Check if an argument was provided that depends on a module that may
  1639.    not be part of the Python standard library.
  1640.  
  1641.    If such an argument is supplied, and the module does not exist, exit
  1642.    with an error stating which module is missing.
  1643.    """
  1644.     optional_args = {
  1645.         'json': ('json/simplejson python module', json),
  1646.         'secure': ('SSL support', HTTPSConnection),
  1647.     }
  1648.  
  1649.     for arg, info in optional_args.items():
  1650.         if getattr(args, arg, False) and info[1] is None:
  1651.             raise SystemExit('%s is not installed. --%s is '
  1652.                              'unavailable' % (info[0], arg))
  1653.  
  1654.  
  1655. def printer(string, quiet=False, debug=False, error=False, **kwargs):
  1656.     """Helper function print a string with various features"""
  1657.  
  1658.     if debug and not DEBUG:
  1659.         return
  1660.  
  1661.     if debug:
  1662.         if sys.stdout.isatty():
  1663.             out = '\033[1;30mDEBUG: %s\033[0m' % string
  1664.         else:
  1665.             out = 'DEBUG: %s' % string
  1666.     else:
  1667.         out = string
  1668.  
  1669.     if error:
  1670.         kwargs['file'] = sys.stderr
  1671.  
  1672.     if not quiet:
  1673.         print_(out, **kwargs)
  1674.  
  1675.  
  1676. def shell():
  1677.     """Run the full speedtest.net test"""
  1678.  
  1679.     global DEBUG
  1680.     shutdown_event = threading.Event()
  1681.  
  1682.     signal.signal(signal.SIGINT, ctrl_c(shutdown_event))
  1683.  
  1684.     args = parse_args()
  1685.  
  1686.     # Print the version and exit
  1687.     if args.version:
  1688.         version()
  1689.  
  1690.     if not args.download and not args.upload:
  1691.         raise SpeedtestCLIError('Cannot supply both --no-download and '
  1692.                                 '--no-upload')
  1693.  
  1694.     if len(args.csv_delimiter) != 1:
  1695.         raise SpeedtestCLIError('--csv-delimiter must be a single character')
  1696.  
  1697.     if args.csv_header:
  1698.         csv_header(args.csv_delimiter)
  1699.  
  1700.     validate_optional_args(args)
  1701.  
  1702.     debug = getattr(args, 'debug', False)
  1703.     if debug == 'SUPPRESSHELP':
  1704.         debug = False
  1705.     if debug:
  1706.         DEBUG = True
  1707.  
  1708.     if args.simple or args.csv or args.json:
  1709.         quiet = True
  1710.     else:
  1711.         quiet = False
  1712.  
  1713.     if args.csv or args.json:
  1714.         machine_format = True
  1715.     else:
  1716.         machine_format = False
  1717.  
  1718.     # Don't set a callback if we are running quietly
  1719.     if quiet or debug:
  1720.         callback = do_nothing
  1721.     else:
  1722.         callback = print_dots(shutdown_event)
  1723.  
  1724.     printer('Retrieving speedtest.net configuration...', quiet)
  1725.     try:
  1726.         speedtest = Speedtest(
  1727.             source_address=args.source,
  1728.             timeout=args.timeout,
  1729.             secure=args.secure
  1730.         )
  1731.     except (ConfigRetrievalError,) + HTTP_ERRORS:
  1732.         printer('Cannot retrieve speedtest configuration', error=True)
  1733.         raise SpeedtestCLIError(get_exception())
  1734.  
  1735.     if args.list:
  1736.         try:
  1737.             speedtest.get_servers()
  1738.         except (ServersRetrievalError,) + HTTP_ERRORS:
  1739.             printer('Cannot retrieve speedtest server list', error=True)
  1740.             raise SpeedtestCLIError(get_exception())
  1741.  
  1742.         for _, servers in sorted(speedtest.servers.items()):
  1743.             for server in servers:
  1744.                 line = ('%(id)5s) %(sponsor)s (%(name)s, %(country)s) '
  1745.                         '[%(d)0.2f km]' % server)
  1746.                 try:
  1747.                     printer(line)
  1748.                 except IOError:
  1749.                     e = get_exception()
  1750.                     if e.errno != errno.EPIPE:
  1751.                         raise
  1752.         sys.exit(0)
  1753.  
  1754.     printer('Testing from %(isp)s (%(ip)s)...' % speedtest.config['client'],
  1755.             quiet)
  1756.  
  1757.     if not args.mini:
  1758.         printer('Retrieving speedtest.net server list...', quiet)
  1759.         try:
  1760.             speedtest.get_servers(servers=args.server, exclude=args.exclude)
  1761.         except NoMatchedServers:
  1762.             raise SpeedtestCLIError(
  1763.                 'No matched servers: %s' %
  1764.                 ', '.join('%s' % s for s in args.server)
  1765.             )
  1766.         except (ServersRetrievalError,) + HTTP_ERRORS:
  1767.             printer('Cannot retrieve speedtest server list', error=True)
  1768.             raise SpeedtestCLIError(get_exception())
  1769.         except InvalidServerIDType:
  1770.             raise SpeedtestCLIError(
  1771.                 '%s is an invalid server type, must '
  1772.                 'be an int' % ', '.join('%s' % s for s in args.server)
  1773.             )
  1774.  
  1775.         if args.server and len(args.server) == 1:
  1776.             printer('Retrieving information for the selected server...', quiet)
  1777.         else:
  1778.             printer('Selecting best server based on ping...', quiet)
  1779.         speedtest.get_best_server()
  1780.     elif args.mini:
  1781.         speedtest.get_best_server(speedtest.set_mini_server(args.mini))
  1782.  
  1783.     results = speedtest.results
  1784.  
  1785.     printer('Hosted by %(sponsor)s (%(name)s) [%(d)0.2f km]: '
  1786.             '%(latency)s ms' % results.server, quiet)
  1787.  
  1788.     if args.download:
  1789.         printer('Testing download speed', quiet,
  1790.                 end=('', '\n')[bool(debug)])
  1791.         speedtest.download(callback=callback)
  1792.         printer('Download: %0.2f M%s/s' %
  1793.                 ((results.download / 1000.0 / 1000.0) / args.units[1],
  1794.                  args.units[0]),
  1795.                 quiet)
  1796.     else:
  1797.         printer('Skipping download test', quiet)
  1798.  
  1799.     if args.upload:
  1800.         printer('Testing upload speed', quiet,
  1801.                 end=('', '\n')[bool(debug)])
  1802.         speedtest.upload(callback=callback, pre_allocate=args.pre_allocate)
  1803.         printer('Upload: %0.2f M%s/s' %
  1804.                 ((results.upload / 1000.0 / 1000.0) / args.units[1],
  1805.                  args.units[0]),
  1806.                 quiet)
  1807.     else:
  1808.         printer('Skipping upload test', quiet)
  1809.  
  1810.     printer('Results:\n%r' % results.dict(), debug=True)
  1811.  
  1812.     if args.simple:
  1813.         printer('Ping: %s ms\nDownload: %0.2f M%s/s\nUpload: %0.2f M%s/s' %
  1814.                 (results.ping,
  1815.                  (results.download / 1000.0 / 1000.0) / args.units[1],
  1816.                  args.units[0],
  1817.                  (results.upload / 1000.0 / 1000.0) / args.units[1],
  1818.                  args.units[0]))
  1819.     elif args.csv:
  1820.         printer(results.csv(delimiter=args.csv_delimiter))
  1821.     elif args.json:
  1822.         if args.share:
  1823.             results.share()
  1824.         printer(results.json())
  1825.  
  1826.     if args.share and not machine_format:
  1827.         printer('Share results: %s' % results.share())
  1828.  
  1829.  
  1830. def main():
  1831.     try:
  1832.         shell()
  1833.     except KeyboardInterrupt:
  1834.         printer('\nCancelling...', error=True)
  1835.     except (SpeedtestException, SystemExit):
  1836.         e = get_exception()
  1837.         # Ignore a successful exit, or argparse exit
  1838.         if getattr(e, 'code', 1) not in (0, 2):
  1839.             raise SystemExit('ERROR: %s' % e)
  1840.  
  1841.  
  1842. if __name__ == '__main__':
  1843.     main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement