server.py 6.38 KB
Newer Older
Matt Traudt's avatar
Matt Traudt committed
1
from ..util.simpleauth import authenticate_client
Matt Traudt's avatar
Matt Traudt committed
2
from ..util.sockio import read_line
3
from sbws.globals import (fail_hard, is_initted)
4
from sbws.globals import (MIN_REQ_BYTES, MAX_REQ_BYTES, SOCKET_TIMEOUT)
Matt Traudt's avatar
Matt Traudt committed
5
from argparse import ArgumentDefaultsHelpFormatter
6
from functools import lru_cache
Matt Traudt's avatar
Matt Traudt committed
7
from threading import Thread
Matt Traudt's avatar
Matt Traudt committed
8
import socket
9
import time
10
import random
11
import os
12
13


Matt Traudt's avatar
Matt Traudt committed
14
def gen_parser(sub):
15
16
17
18
19
    d = 'The server side of sbws. This should be run on the same machine as '\
        'a helper relay. This listens for clients connections and responds '\
        'with the number of bytes the client requests.'
    sub.add_parser('server', formatter_class=ArgumentDefaultsHelpFormatter,
                   description=d)
Matt Traudt's avatar
Matt Traudt committed
20
21


22
def close_socket(s):
23
    try:
24
        log.info('Closing fd', s.fileno())
25
26
27
28
        s.shutdown(socket.SHUT_RDWR)
        s.close()
    except OSError:
        pass
29
30
31


def get_send_amount(sock):
Matt Traudt's avatar
Matt Traudt committed
32
    line = read_line(sock, max_len=16, log_fn=log.info)
Matt Traudt's avatar
Matt Traudt committed
33
34
    if line is None:
        return None
35
36
37
38
39
    # if len(line) == 16, then it is much more likely we read garbage or not an
    # entire line instead of a legit number of bytes to send. So say we've
    # failed.
    if len(line) == 16:
        return None
40
    try:
Matt Traudt's avatar
Matt Traudt committed
41
        send_amount = int(line)
42
43
44
45
46
    except (TypeError, ValueError):
        return None
    return send_amount


47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
@lru_cache(maxsize=8)
def _generate_random_string(length):
    ''' Generates a VERY WEAKLY random string. It felt wrong only sending a
    ton of a single character, but generating a long and "truely" random
    string is way too expensive. Furthermore, we don't just send random bytes
    (which may be easy to generate) because for some reason I have it in my
    head that doing everything in simple ascii, 1 byte == 1 char, and sometimes
    line-based way is a smart idea.

    Anyway. This shuffles the alphabet. It then concatenates this shuffled
    alphabet as many times as necessary to get a string as long or longer than
    the required length. It then returns the string up until the required
    length.

    Oh. Also it caches a few results based on the requested length. That's
    another thing that hurts its randomness.
    '''
    assert length > 0
65
    # start = time_now()
66
    repeats = int(length / len(_generate_random_string.alphabet)) + 1
67
    rng.shuffle(_generate_random_string.alphabet)
68
69
    s = ''.join(_generate_random_string.alphabet)
    s = s * repeats
70
    # stop = time_now()
71
72
73
74
75
76
77
78
79
80
81
82
83
84
    # _generate_random_string.acc += stop - start
    # if stop >= 60 + _generate_random_string.last_log:
    #     log.notice('Spent', _generate_random_string.acc,
    #                'seconds in the last minute generating "random" strings')
    #     _generate_random_string.acc = 0
    #     _generate_random_string.last_log = stop
    assert len(s) >= length
    return s[:length]


_generate_random_string.alphabet = list('abcdefghijklmnopqrstuvwxyz'
                                        'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
                                        '0123456789')
# _generate_random_string.acc = 0
85
# _generate_random_string.last_log = time_now()
86
87


88
def write_to_client(sock, conf, amount):
89
    ''' Returns True if successful; else False '''
90
    log.info('Sending client no.', sock.fileno(), amount, 'bytes')
91
    while amount > 0:
92
93
        amount_this_time = min(conf.getint('server', 'max_send_per_write'),
                               amount)
94
95
        amount -= amount_this_time
        try:
96
97
            sock.send(bytes(
                _generate_random_string(amount_this_time), 'utf-8'))
98
        except (socket.timeout, ConnectionResetError, BrokenPipeError) as e:
99
            log.info('fd', sock.fileno(), ':', e)
100
101
            return False
    return True
102
103


104
def new_thread(args, conf, sock):
Matt Traudt's avatar
Matt Traudt committed
105
    def closure():
106
107
108
        client_name = authenticate_client(
            sock, conf['server.passwords'], log.info)
        if not client_name:
109
110
111
            log.info('Client did not provide valid auth')
            close_socket(sock)
            return
112
        log.notice(client_name, 'authenticated on', sock.fileno())
113
114
115
        while True:
            send_amount = get_send_amount(sock)
            if send_amount is None:
116
                log.info('Couldn\'t get an amount to send to', sock.fileno())
Matt Traudt's avatar
Matt Traudt committed
117
                break
118
119
120
121
            if send_amount < MIN_REQ_BYTES or send_amount > MAX_REQ_BYTES:
                log.warn(client_name, 'requested', send_amount, 'bytes, which '
                         'is not valid')
                break
122
            write_to_client(sock, conf, send_amount)
Matt Traudt's avatar
Matt Traudt committed
123
        log.notice(client_name, 'on', sock.fileno(), 'went away')
124
        close_socket(sock)
Matt Traudt's avatar
Matt Traudt committed
125
126
127
    thread = Thread(target=closure)
    return thread

128

129
def main(args, conf, log_):
Matt Traudt's avatar
Matt Traudt committed
130
    global log
131
    global rng
Matt Traudt's avatar
Matt Traudt committed
132
    log = log_
133
    rng = random.SystemRandom()
134
    if not is_initted(args.directory):
135
        fail_hard('Sbws isn\'t initialized. Try sbws init', log=log)
Matt Traudt's avatar
Matt Traudt committed
136

137
    if len(conf['server.passwords']) < 1:
138
        conf_fname = os.path.join(args.directory, 'config.ini')
Matt Traudt's avatar
Matt Traudt committed
139
140
141
        fail_hard('Sbws server needs at least one password in the section'
                  ' [server.passwords] in the config file in {}. See '
                  'DEPLOY.rst for more information.'
142
                  .format(conf_fname), log=log)
Matt Traudt's avatar
Matt Traudt committed
143

144
    h = (conf['server']['bind_ip'], conf.getint('server', 'bind_port'))
145
    log.notice('Binding to', h)
146
    while True:
Matt Traudt's avatar
Matt Traudt committed
147
        try:
148
149
150
            # first try IPv4
            log.debug('Trying to bind while assuming ipv4')
            server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
Matt Traudt's avatar
Matt Traudt committed
151
            server.bind(h)
152
153
154
155
156
157
158
159
160
161
162
163
        except OSError as e1:
            try:
                # then try IPv6
                log.debug('Trying to bind while assuming ipv6')
                server = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
                server.bind(h)
            except OSError as e2:
                log.warn(e1)
                log.warn(e2)
                time.sleep(5)
            else:
                break
Matt Traudt's avatar
Matt Traudt committed
164
165
        else:
            break
166
    log.notice('Listening on', h)
Matt Traudt's avatar
Matt Traudt committed
167
168
169
170
    server.listen(5)
    try:
        while True:
            sock, addr = server.accept()
171
            sock.settimeout(SOCKET_TIMEOUT)
172
            log.info('accepting connection from', addr, 'as', sock.fileno())
173
            t = new_thread(args, conf, sock)
174
            t.start()
Matt Traudt's avatar
Matt Traudt committed
175
176
177
    except KeyboardInterrupt:
        pass
    finally:
178
179
        log.info('Generate random string stats:',
                 _generate_random_string.cache_info())
180
        close_socket(server)