Advertisement
tirabytes

Python: Speedtest

Apr 28th, 2014
515
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 20.21 KB | None | 0 0
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. # Copyright 2013 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. __version__ = '0.2.5'
  19.  
  20. # Some global variables we use
  21. source = None
  22. shutdown_event = None
  23.  
  24. import math
  25. import time
  26. import os
  27. import sys
  28. import threading
  29. import re
  30. import signal
  31. import socket
  32.  
  33. # Used for bound_interface
  34. socket_socket = socket.socket
  35.  
  36. try:
  37.     import xml.etree.cElementTree as ET
  38. except ImportError:
  39.     try:
  40.         import xml.etree.ElementTree as ET
  41.     except ImportError:
  42.         from xml.dom import minidom as DOM
  43.         ET = None
  44.  
  45. # Begin import game to handle Python 2 and Python 3
  46. try:
  47.     from urllib2 import urlopen, Request, HTTPError, URLError
  48. except ImportError:
  49.     from urllib.request import urlopen, Request, HTTPError, URLError
  50.  
  51. try:
  52.     from Queue import Queue
  53. except ImportError:
  54.     from queue import Queue
  55.  
  56. try:
  57.     from urlparse import urlparse
  58. except ImportError:
  59.     from urllib.parse import urlparse
  60.  
  61. try:
  62.     from urlparse import parse_qs
  63. except ImportError:
  64.     try:
  65.         from urllib.parse import parse_qs
  66.     except ImportError:
  67.         from cgi import parse_qs
  68.  
  69. try:
  70.     from hashlib import md5
  71. except ImportError:
  72.     from md5 import md5
  73.  
  74. try:
  75.     from argparse import ArgumentParser as ArgParser
  76. except ImportError:
  77.     from optparse import OptionParser as ArgParser
  78.  
  79. try:
  80.     import builtins
  81. except ImportError:
  82.     def print_(*args, **kwargs):
  83.         """The new-style print function taken from
  84.        https://pypi.python.org/pypi/six/
  85.  
  86.        """
  87.         fp = kwargs.pop("file", sys.stdout)
  88.         if fp is None:
  89.             return
  90.  
  91.         def write(data):
  92.             if not isinstance(data, basestring):
  93.                 data = str(data)
  94.             fp.write(data)
  95.  
  96.         want_unicode = False
  97.         sep = kwargs.pop("sep", None)
  98.         if sep is not None:
  99.             if isinstance(sep, unicode):
  100.                 want_unicode = True
  101.             elif not isinstance(sep, str):
  102.                 raise TypeError("sep must be None or a string")
  103.         end = kwargs.pop("end", None)
  104.         if end is not None:
  105.             if isinstance(end, unicode):
  106.                 want_unicode = True
  107.             elif not isinstance(end, str):
  108.                 raise TypeError("end must be None or a string")
  109.         if kwargs:
  110.             raise TypeError("invalid keyword arguments to print()")
  111.         if not want_unicode:
  112.             for arg in args:
  113.                 if isinstance(arg, unicode):
  114.                     want_unicode = True
  115.                     break
  116.         if want_unicode:
  117.             newline = unicode("\n")
  118.             space = unicode(" ")
  119.         else:
  120.             newline = "\n"
  121.             space = " "
  122.         if sep is None:
  123.             sep = space
  124.         if end is None:
  125.             end = newline
  126.         for i, arg in enumerate(args):
  127.             if i:
  128.                 write(sep)
  129.             write(arg)
  130.         write(end)
  131. else:
  132.     print_ = getattr(builtins, 'print')
  133.     del builtins
  134.  
  135.  
  136. def bound_socket(*args, **kwargs):
  137.     """Bind socket to a specified source IP address"""
  138.  
  139.     global source
  140.     sock = socket_socket(*args, **kwargs)
  141.     sock.bind((source, 0))
  142.     return sock
  143.  
  144.  
  145. def distance(origin, destination):
  146.     """Determine distance between 2 sets of [lat,lon] in km"""
  147.  
  148.     lat1, lon1 = origin
  149.     lat2, lon2 = destination
  150.     radius = 6371  # km
  151.  
  152.     dlat = math.radians(lat2 - lat1)
  153.     dlon = math.radians(lon2 - lon1)
  154.     a = (math.sin(dlat / 2) * math.sin(dlat / 2) + math.cos(math.radians(lat1))
  155.          * math.cos(math.radians(lat2)) * math.sin(dlon / 2)
  156.          * math.sin(dlon / 2))
  157.     c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
  158.     d = radius * c
  159.  
  160.     return d
  161.  
  162.  
  163. class FileGetter(threading.Thread):
  164.     """Thread class for retrieving a URL"""
  165.  
  166.     def __init__(self, url, start):
  167.         self.url = url
  168.         self.result = None
  169.         self.starttime = start
  170.         threading.Thread.__init__(self)
  171.  
  172.     def run(self):
  173.         self.result = [0]
  174.         try:
  175.             if (time.time() - self.starttime) <= 10:
  176.                 f = urlopen(self.url)
  177.                 while 1 and not shutdown_event.isSet():
  178.                     self.result.append(len(f.read(10240)))
  179.                     if self.result[-1] == 0:
  180.                         break
  181.                 f.close()
  182.         except IOError:
  183.             pass
  184.  
  185.  
  186. def downloadSpeed(files, quiet=False):
  187.     """Function to launch FileGetter threads and calculate download speeds"""
  188.  
  189.     start = time.time()
  190.  
  191.     def producer(q, files):
  192.         for file in files:
  193.             thread = FileGetter(file, start)
  194.             thread.start()
  195.             q.put(thread, True)
  196.             if not quiet and not shutdown_event.isSet():
  197.                 sys.stdout.write('.')
  198.                 sys.stdout.flush()
  199.  
  200.     finished = []
  201.  
  202.     def consumer(q, total_files):
  203.         while len(finished) < total_files:
  204.             thread = q.get(True)
  205.             while thread.isAlive():
  206.                 thread.join(timeout=0.1)
  207.             finished.append(sum(thread.result))
  208.             del thread
  209.  
  210.     q = Queue(6)
  211.     prod_thread = threading.Thread(target=producer, args=(q, files))
  212.     cons_thread = threading.Thread(target=consumer, args=(q, len(files)))
  213.     start = time.time()
  214.     prod_thread.start()
  215.     cons_thread.start()
  216.     while prod_thread.isAlive():
  217.         prod_thread.join(timeout=0.1)
  218.     while cons_thread.isAlive():
  219.         cons_thread.join(timeout=0.1)
  220.     return (sum(finished) / (time.time() - start))
  221.  
  222.  
  223. class FilePutter(threading.Thread):
  224.     """Thread class for putting a URL"""
  225.  
  226.     def __init__(self, url, start, size):
  227.         self.url = url
  228.         chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ'
  229.         data = chars * (int(round(int(size) / 36.0)))
  230.         self.data = ('content1=%s' % data[0:int(size) - 9]).encode()
  231.         del data
  232.         self.result = None
  233.         self.starttime = start
  234.         threading.Thread.__init__(self)
  235.  
  236.     def run(self):
  237.         try:
  238.             if ((time.time() - self.starttime) <= 10 and
  239.                     not shutdown_event.isSet()):
  240.                 f = urlopen(self.url, self.data)
  241.                 f.read(11)
  242.                 f.close()
  243.                 self.result = len(self.data)
  244.             else:
  245.                 self.result = 0
  246.         except IOError:
  247.             self.result = 0
  248.  
  249.  
  250. def uploadSpeed(url, sizes, quiet=False):
  251.     """Function to launch FilePutter threads and calculate upload speeds"""
  252.  
  253.     start = time.time()
  254.  
  255.     def producer(q, sizes):
  256.         for size in sizes:
  257.             thread = FilePutter(url, start, size)
  258.             thread.start()
  259.             q.put(thread, True)
  260.             if not quiet and not shutdown_event.isSet():
  261.                 sys.stdout.write('.')
  262.                 sys.stdout.flush()
  263.  
  264.     finished = []
  265.  
  266.     def consumer(q, total_sizes):
  267.         while len(finished) < total_sizes:
  268.             thread = q.get(True)
  269.             while thread.isAlive():
  270.                 thread.join(timeout=0.1)
  271.             finished.append(thread.result)
  272.             del thread
  273.  
  274.     q = Queue(6)
  275.     prod_thread = threading.Thread(target=producer, args=(q, sizes))
  276.     cons_thread = threading.Thread(target=consumer, args=(q, len(sizes)))
  277.     start = time.time()
  278.     prod_thread.start()
  279.     cons_thread.start()
  280.     while prod_thread.isAlive():
  281.         prod_thread.join(timeout=0.1)
  282.     while cons_thread.isAlive():
  283.         cons_thread.join(timeout=0.1)
  284.     return (sum(finished) / (time.time() - start))
  285.  
  286.  
  287. def getAttributesByTagName(dom, tagName):
  288.     """Retrieve an attribute from an XML document and return it in a
  289.    consistent format
  290.  
  291.    Only used with xml.dom.minidom, which is likely only to be used
  292.    with python versions older than 2.5
  293.    """
  294.     elem = dom.getElementsByTagName(tagName)[0]
  295.     return dict(list(elem.attributes.items()))
  296.  
  297.  
  298. def getConfig():
  299.     """Download the speedtest.net configuration and return only the data
  300.    we are interested in
  301.    """
  302.  
  303.     uh = urlopen('http://www.speedtest.net/speedtest-config.php')
  304.     configxml = []
  305.     while 1:
  306.         configxml.append(uh.read(10240))
  307.         if len(configxml[-1]) == 0:
  308.             break
  309.     if int(uh.code) != 200:
  310.         return None
  311.     uh.close()
  312.     try:
  313.         root = ET.fromstring(''.encode().join(configxml))
  314.         config = {
  315.             'client': root.find('client').attrib,
  316.             'times': root.find('times').attrib,
  317.             'download': root.find('download').attrib,
  318.             'upload': root.find('upload').attrib}
  319.     except AttributeError:
  320.         root = DOM.parseString(''.join(configxml))
  321.         config = {
  322.             'client': getAttributesByTagName(root, 'client'),
  323.             'times': getAttributesByTagName(root, 'times'),
  324.             'download': getAttributesByTagName(root, 'download'),
  325.             'upload': getAttributesByTagName(root, 'upload')}
  326.     del root
  327.     del configxml
  328.     return config
  329.  
  330.  
  331. def closestServers(client, all=False):
  332.     """Determine the 5 closest speedtest.net servers based on geographic
  333.    distance
  334.    """
  335.  
  336.     uh = urlopen('http://www.speedtest.net/speedtest-servers.php')
  337.     serversxml = []
  338.     while 1:
  339.         serversxml.append(uh.read(10240))
  340.         if len(serversxml[-1]) == 0:
  341.             break
  342.     if int(uh.code) != 200:
  343.         return None
  344.     uh.close()
  345.     try:
  346.         root = ET.fromstring(''.encode().join(serversxml))
  347.         elements = root.getiterator('server')
  348.     except AttributeError:
  349.         root = DOM.parseString(''.join(serversxml))
  350.         elements = root.getElementsByTagName('server')
  351.     servers = {}
  352.     for server in elements:
  353.         try:
  354.             attrib = server.attrib
  355.         except AttributeError:
  356.             attrib = dict(list(server.attributes.items()))
  357.         d = distance([float(client['lat']), float(client['lon'])],
  358.                      [float(attrib.get('lat')), float(attrib.get('lon'))])
  359.         attrib['d'] = d
  360.         if d not in servers:
  361.             servers[d] = [attrib]
  362.         else:
  363.             servers[d].append(attrib)
  364.     del root
  365.     del serversxml
  366.     del elements
  367.  
  368.     closest = []
  369.     for d in sorted(servers.keys()):
  370.         for s in servers[d]:
  371.             closest.append(s)
  372.             if len(closest) == 5 and not all:
  373.                 break
  374.         else:
  375.             continue
  376.         break
  377.  
  378.     del servers
  379.     return closest
  380.  
  381.  
  382. def getBestServer(servers):
  383.     """Perform a speedtest.net "ping" to determine which speedtest.net
  384.    server has the lowest latency
  385.    """
  386.  
  387.     results = {}
  388.     for server in servers:
  389.         cum = []
  390.         url = os.path.dirname(server['url'])
  391.         for i in range(0, 3):
  392.             try:
  393.                 uh = urlopen('%s/latency.txt' % url)
  394.             except (HTTPError, URLError):
  395.                 cum.append(3600)
  396.                 continue
  397.             start = time.time()
  398.             text = uh.read(9)
  399.             total = time.time() - start
  400.             if int(uh.code) == 200 and text == 'test=test'.encode():
  401.                 cum.append(total)
  402.             else:
  403.                 cum.append(3600)
  404.             uh.close()
  405.         avg = round((sum(cum) / 3) * 1000000, 3)
  406.         results[avg] = server
  407.  
  408.     fastest = sorted(results.keys())[0]
  409.     best = results[fastest]
  410.     best['latency'] = fastest
  411.  
  412.     return best
  413.  
  414.  
  415. def ctrl_c(signum, frame):
  416.     """Catch Ctrl-C key sequence and set a shutdown_event for our threaded
  417.    operations
  418.    """
  419.  
  420.     global shutdown_event
  421.     shutdown_event.set()
  422.     raise SystemExit('\nCancelling...')
  423.  
  424.  
  425. def version():
  426.     """Print the version"""
  427.  
  428.     raise SystemExit(__version__)
  429.  
  430.  
  431. def speedtest():
  432.     """Run the full speedtest.net test"""
  433.  
  434.     global shutdown_event, source
  435.     shutdown_event = threading.Event()
  436.  
  437.     signal.signal(signal.SIGINT, ctrl_c)
  438.  
  439.     description = (
  440.         'Command line interface for testing internet bandwidth using '
  441.         'speedtest.net.\n'
  442.         '------------------------------------------------------------'
  443.         '--------------\n'
  444.         'https://github.com/sivel/speedtest-cli')
  445.  
  446.     parser = ArgParser(description=description)
  447.     # Give optparse.OptionParser an `add_argument` method for
  448.     # compatibility with argparse.ArgumentParser
  449.     try:
  450.         parser.add_argument = parser.add_option
  451.     except AttributeError:
  452.         pass
  453.     parser.add_argument('--share', action='store_true',
  454.                         help='Generate and provide a URL to the speedtest.net '
  455.                              'share results image')
  456.     parser.add_argument('--simple', action='store_true',
  457.                         help='Suppress verbose output, only show basic '
  458.                              'information')
  459.     parser.add_argument('--list', action='store_true',
  460.                         help='Display a list of speedtest.net servers '
  461.                              'sorted by distance')
  462.     parser.add_argument('--server', help='Specify a server ID to test against')
  463.     parser.add_argument('--mini', help='URL of the Speedtest Mini server')
  464.     parser.add_argument('--source', help='Source IP address to bind to')
  465.     parser.add_argument('--version', action='store_true',
  466.                         help='Show the version number and exit')
  467.  
  468.     options = parser.parse_args()
  469.     if isinstance(options, tuple):
  470.         args = options[0]
  471.     else:
  472.         args = options
  473.     del options
  474.  
  475.     # Print the version and exit
  476.     if args.version:
  477.         version()
  478.  
  479.     # If specified bind to a specific IP address
  480.     if args.source:
  481.         source = args.source
  482.         socket.socket = bound_socket
  483.  
  484.     if not args.simple:
  485.         print_('Retrieving speedtest.net configuration...')
  486.     try:
  487.         config = getConfig()
  488.     except URLError:
  489.         print_('Cannot retrieve speedtest configuration')
  490.         sys.exit(1)
  491.  
  492.     if not args.simple:
  493.         print_('Retrieving speedtest.net server list...')
  494.     if args.list or args.server:
  495.         servers = closestServers(config['client'], True)
  496.         if args.list:
  497.             serverList = []
  498.             for server in servers:
  499.                 line = ('%(id)4s) %(sponsor)s (%(name)s, %(country)s) '
  500.                         '[%(d)0.2f km]' % server)
  501.                 serverList.append(line)
  502.             # Python 2.7 and newer seem to be ok with the resultant encoding
  503.             # from parsing the XML, but older versions have some issues.
  504.             # This block should detect whether we need to encode or not
  505.             try:
  506.                 unicode()
  507.                 print_('\n'.join(serverList).encode('utf-8', 'ignore'))
  508.             except NameError:
  509.                 print_('\n'.join(serverList))
  510.             except IOError:
  511.                 pass
  512.             sys.exit(0)
  513.     else:
  514.         servers = closestServers(config['client'])
  515.  
  516.     if not args.simple:
  517.         print_('Testing from %(isp)s (%(ip)s)...' % config['client'])
  518.  
  519.     if args.server:
  520.         try:
  521.             best = getBestServer(filter(lambda x: x['id'] == args.server,
  522.                                         servers))
  523.         except IndexError:
  524.             print_('Invalid server ID')
  525.             sys.exit(1)
  526.     elif args.mini:
  527.         name, ext = os.path.splitext(args.mini)
  528.         if ext:
  529.             url = os.path.dirname(args.mini)
  530.         else:
  531.             url = args.mini
  532.         urlparts = urlparse(url)
  533.         try:
  534.             f = urlopen(args.mini)
  535.         except:
  536.             print_('Invalid Speedtest Mini URL')
  537.             sys.exit(1)
  538.         else:
  539.             text = f.read()
  540.             f.close()
  541.         extension = re.findall('upload_extension: "([^"]+)"', text.decode())
  542.         if not urlparts or not extension:
  543.             print_('Please provide the full URL of your Speedtest Mini server')
  544.             sys.exit(1)
  545.         servers = [{
  546.             'sponsor': 'Speedtest Mini',
  547.             'name': urlparts[1],
  548.             'd': 0,
  549.             'url': '%s/speedtest/upload.%s' % (url.rstrip('/'), extension[0]),
  550.             'latency': 0,
  551.             'id': 0
  552.         }]
  553.         try:
  554.             best = getBestServer(servers)
  555.         except:
  556.             best = servers[0]
  557.     else:
  558.         if not args.simple:
  559.             print_('Selecting best server based on ping...')
  560.         best = getBestServer(servers)
  561.  
  562.     if not args.simple:
  563.         # Python 2.7 and newer seem to be ok with the resultant encoding
  564.         # from parsing the XML, but older versions have some issues.
  565.         # This block should detect whether we need to encode or not
  566.         try:
  567.             unicode()
  568.             print_(('Hosted by %(sponsor)s (%(name)s) [%(d)0.2f km]: '
  569.                    '%(latency)s ms' % best).encode('utf-8', 'ignore'))
  570.         except NameError:
  571.             print_('Hosted by %(sponsor)s (%(name)s) [%(d)0.2f km]: '
  572.                    '%(latency)s ms' % best)
  573.     else:
  574.         print_('Ping: %(latency)s ms' % best)
  575.  
  576.     sizes = [350, 500, 750, 1000, 1500, 2000, 2500, 3000, 3500, 4000]
  577.     urls = []
  578.     for size in sizes:
  579.         for i in range(0, 4):
  580.             urls.append('%s/random%sx%s.jpg' %
  581.                         (os.path.dirname(best['url']), size, size))
  582.     if not args.simple:
  583.         print_('Testing download speed', end='')
  584.     dlspeed = downloadSpeed(urls, args.simple)
  585.     if not args.simple:
  586.         print_()
  587.     print_('Download: %0.2f Mbit/s' % ((dlspeed / 1000 / 1000) * 8))
  588.  
  589.     sizesizes = [int(.25 * 1000 * 1000), int(.5 * 1000 * 1000)]
  590.     sizes = []
  591.     for size in sizesizes:
  592.         for i in range(0, 25):
  593.             sizes.append(size)
  594.     if not args.simple:
  595.         print_('Testing upload speed', end='')
  596.     ulspeed = uploadSpeed(best['url'], sizes, args.simple)
  597.     if not args.simple:
  598.         print_()
  599.     print_('Upload: %0.2f Mbit/s' % ((ulspeed / 1000 / 1000) * 8))
  600.  
  601.     if args.share and args.mini:
  602.         print_('Cannot generate a speedtest.net share results image while '
  603.                'testing against a Speedtest Mini server')
  604.     elif args.share:
  605.         dlspeedk = int(round((dlspeed / 1000) * 8, 0))
  606.         ping = int(round(best['latency'], 0))
  607.         ulspeedk = int(round((ulspeed / 1000) * 8, 0))
  608.  
  609.         # Build the request to send results back to speedtest.net
  610.         # We use a list instead of a dict because the API expects parameters
  611.         # in a certain order
  612.         apiData = [
  613.             'download=%s' % dlspeedk,
  614.             'ping=%s' % ping,
  615.             'upload=%s' % ulspeedk,
  616.             'promo=',
  617.             'startmode=%s' % 'pingselect',
  618.             'recommendedserverid=%s' % best['id'],
  619.             'accuracy=%s' % 1,
  620.             'serverid=%s' % best['id'],
  621.             'hash=%s' % md5(('%s-%s-%s-%s' %
  622.                              (ping, ulspeedk, dlspeedk, '297aae72'))
  623.                             .encode()).hexdigest()]
  624.  
  625.         req = Request('http://www.speedtest.net/api/api.php',
  626.                       data='&'.join(apiData).encode())
  627.         req.add_header('Referer', 'http://c.speedtest.net/flash/speedtest.swf')
  628.         f = urlopen(req)
  629.         response = f.read()
  630.         code = f.code
  631.         f.close()
  632.  
  633.         if int(code) != 200:
  634.             print_('Could not submit results to speedtest.net')
  635.             sys.exit(1)
  636.  
  637.         qsargs = parse_qs(response.decode())
  638.         resultid = qsargs.get('resultid')
  639.         if not resultid or len(resultid) != 1:
  640.             print_('Could not submit results to speedtest.net')
  641.             sys.exit(1)
  642.  
  643.         print_('Share results: http://www.speedtest.net/result/%s.png' %
  644.                resultid[0])
  645.  
  646.  
  647. def main():
  648.     try:
  649.         speedtest()
  650.     except KeyboardInterrupt:
  651.         print_('\nCancelling...')
  652.  
  653.  
  654. if __name__ == '__main__':
  655.     main()
  656.  
  657. # vim:ts=4:sw=4:expandtab
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement