voussoir/else
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add overall_ratelimit, individual_ratelimit.

Browse source
This commit is contained in:
voussoir 2021-02-20 13:05:45 -08:00
parent 4b599e38d3
commit 0dec7cbcff
No known key found for this signature in database
GPG key ID: 5F7554F8C26DACCB
1 changed files with 42 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
SimpleServer/simpleserver.py
View file

@ -3,6 +3,7 @@ import base64
import cgi import cgi
import http.cookies import http.cookies
import http.server import http.server
import math
import mimetypes import mimetypes
import os import os
import sys import sys
@ -12,12 +13,12 @@ import zipstream
from voussoirkit import betterhelp from voussoirkit import betterhelp
from voussoirkit import bytestring from voussoirkit import bytestring
from voussoirkit import gentools
from voussoirkit import passwordy from voussoirkit import passwordy
from voussoirkit import pathclass from voussoirkit import pathclass
from voussoirkit import ratelimiter from voussoirkit import ratelimiter
FILE_READ_CHUNK = bytestring.MIBIBYTE CHUNK_SIZE = bytestring.MIBIBYTE
RATELIMITER = ratelimiter.Ratelimiter(16 * bytestring.MIBIBYTE)
OPENDIR_TEMPLATE = ''' OPENDIR_TEMPLATE = '''
<html> <html>
@ -56,8 +57,9 @@ TOKEN_COOKIE_NAME = 'simpleserver_token'
# SERVER ########################################################################################### # SERVER ###########################################################################################
class RequestHandler(http.server.BaseHTTPRequestHandler): class RequestHandler(http.server.BaseHTTPRequestHandler):
def __init__(self, *args, **kwargs): def __init__(self, request, client_info, server, individual_ratelimit):
super().__init__(*args, **kwargs) self.individual_ratelimit = ratelimiter.Ratelimiter(individual_ratelimit)
super().__init__(request, client_info, server)
@property @property
def auth_cookie(self): def auth_cookie(self):
@ -117,12 +119,16 @@ class RequestHandler(http.server.BaseHTTPRequestHandler):
def write(self, data): def write(self, data):
if isinstance(data, str): if isinstance(data, str):
data = data.encode('utf-8') data = data.encode('utf-8')
if isinstance(data, types.GeneratorType):
if isinstance(data, bytes):
databytes = data
data = (databytes[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE] for i in range(math.ceil(len(data)/CHUNK_SIZE)))
for chunk in data: for chunk in data:
self.wfile.write(chunk) self.wfile.write(chunk)
RATELIMITER.limit(len(chunk)) chunksize = len(chunk)
else: self.server.overall_ratelimit.limit(chunksize)
self.wfile.write(data) self.individual_ratelimit.limit(chunksize)
def do_GET(self): def do_GET(self):
if not self.check_has_password(): if not self.check_has_password():
@ -278,11 +284,22 @@ class RequestHandler(http.server.BaseHTTPRequestHandler):
return False return False
class SimpleServer: class SimpleServer:
def __init__(self, port, password, authorize_by_ip, enable_zip): def __init__(
self,
port,
*,
password,
authorize_by_ip,
enable_zip,
overall_ratelimit,
individual_ratelimit,
):
self.port = port self.port = port
self.password = password self.password = password
self.authorize_by_ip = authorize_by_ip self.authorize_by_ip = authorize_by_ip
self.enable_zip = enable_zip self.enable_zip = enable_zip
self.overall_ratelimit = ratelimiter.Ratelimiter(overall_ratelimit)
self.individual_ratelimit = individual_ratelimit
if authorize_by_ip: if authorize_by_ip:
self.accepted_ips = set() self.accepted_ips = set()
@ -292,7 +309,9 @@ class SimpleServer:
self.accepted_ips = None self.accepted_ips = None
def make_request_handler(self, request, client_info, _server): def make_request_handler(self, request, client_info, _server):
return RequestHandler(request, client_info, self) # We ignore the given _server and use self instead because _server will
# be the ThreadingHTTPServer instance.
return RequestHandler(request, client_info, self, individual_ratelimit=self.individual_ratelimit)
def start(self): def start(self):
server = http.server.ThreadingHTTPServer(('0.0.0.0', self.port), self.make_request_handler) server = http.server.ThreadingHTTPServer(('0.0.0.0', self.port), self.make_request_handler)
@ -393,7 +412,7 @@ def read_filebytes(path, range_min=None, range_max=None):
f.seek(range_min) f.seek(range_min)
sent_amount = 0 sent_amount = 0
while sent_amount < range_span: while sent_amount < range_span:
chunk = f.read(FILE_READ_CHUNK) chunk = f.read(CHUNK_SIZE)
if len(chunk) == 0: if len(chunk) == 0:
break break
@ -470,6 +489,12 @@ flags:
--enable_zip: --enable_zip:
Add a 'zip' link to every directory and allow the user to download the Add a 'zip' link to every directory and allow the user to download the
entire directory as a zip file. entire directory as a zip file.
--overall_ratelimit X:
An integer number of bytes, the maximum bytes/sec of the server overall.
--individual_ratelimit X:
An integer number of bytes, the maximum bytes/sec for any single request.
''' '''
def simpleserver_argparse(args): def simpleserver_argparse(args):
@ -478,6 +503,8 @@ def simpleserver_argparse(args):
password=args.password, password=args.password,
authorize_by_ip=args.authorize_by_ip, authorize_by_ip=args.authorize_by_ip,
enable_zip=args.enable_zip, enable_zip=args.enable_zip,
overall_ratelimit=args.overall_ratelimit,
individual_ratelimit=args.individual_ratelimit,
) )
server.start() server.start()
@ -488,6 +515,8 @@ def main(argv):
parser.add_argument('--password', dest='password', default=None) parser.add_argument('--password', dest='password', default=None)
parser.add_argument('--authorize_by_ip', '--authorize-by-ip', dest='authorize_by_ip', action='store_true') parser.add_argument('--authorize_by_ip', '--authorize-by-ip', dest='authorize_by_ip', action='store_true')
parser.add_argument('--enable_zip', '--enable-zip', dest='enable_zip', action='store_true') parser.add_argument('--enable_zip', '--enable-zip', dest='enable_zip', action='store_true')
parser.add_argument('--overall_ratelimit', '--overall-ratelimit', type=bytestring.parsebytes, default=20*bytestring.MIBIBYTE)
parser.add_argument('--individual_ratelimit', '--individual-ratelimit', type=bytestring.parsebytes, default=10*bytestring.MIBIBYTE)
parser.set_defaults(func=simpleserver_argparse) parser.set_defaults(func=simpleserver_argparse)
return betterhelp.single_main(argv, parser, DOCSTRING) return betterhelp.single_main(argv, parser, DOCSTRING)