DownloadServer.py 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. from http.server import HTTPServer, BaseHTTPRequestHandler
  2. import argparse
  3. import time
  4. import subprocess
  5. import sys
  6. import os
  7. import threading
  8. args = None
  9. outerthread = None
  10. barrier = threading.Barrier(2)
  11. class SimpleHTTPRequestHandler(BaseHTTPRequestHandler):
  12. def do_GET(self):
  13. barrier.wait()
  14. self.send_response(200)
  15. self.end_headers()
  16. data = b'D'
  17. if args.speed_limit:
  18. slow_deadline = time.time()+args.limit_duration
  19. while time.time() < slow_deadline:
  20. self.wfile.write(data)
  21. if args.speed_limit:
  22. time.sleep(1.1)
  23. data = data * 100
  24. self.wfile.write(data)
  25. self.close_connection = True
  26. def runServer(fileName):
  27. httpd = HTTPServer(('localhost', 0), SimpleHTTPRequestHandler)
  28. with open(fileName,"w") as f:
  29. f.write('http://localhost:{}/test'.format(httpd.socket.getsockname()[1]))
  30. httpd.handle_request()
  31. os.remove(fileName)
  32. if __name__ == "__main__":
  33. parser = argparse.ArgumentParser()
  34. parser.add_argument('--speed_limit', help='transfer rate limitation', action='store_true',default=False)
  35. parser.add_argument('--limit_duration', help='duration of the transfer rate limitation',default=1, type=float)
  36. parser.add_argument('--file', help='file to write the url to connect to')
  37. parser.add_argument('--subprocess', action='store_true')
  38. args = parser.parse_args()
  39. if not args.subprocess:
  40. subprocess.Popen([sys.executable]+sys.argv+['--subprocess'],stdin=subprocess.DEVNULL, stderr=subprocess.DEVNULL,stdout=subprocess.DEVNULL)
  41. else:
  42. serverThread = threading.Thread(target=runServer,args=(args.file,))
  43. serverThread.daemon = True
  44. serverThread.start()
  45. barrier.wait(60)
  46. serverThread.join(20)