diff --git a/anonstream/helpers/user.py b/anonstream/helpers/user.py index b4d8346..7bc6e63 100644 --- a/anonstream/helpers/user.py +++ b/anonstream/helpers/user.py @@ -23,7 +23,8 @@ def generate_token_hash_and_tag(token): return token_hash, tag def generate_user( - timestamp, token, broadcaster, verified=False, presence=Presence.NOTWATCHING, + timestamp, token, broadcaster, + verified=False, presence=Presence.NOTWATCHING, headers=None, ): colour = generate_colour( seed='name\0' + token, @@ -54,6 +55,7 @@ def generate_user( 'total': 0, 'current': {}, }, + 'headers': headers, } def get_default_name(user): diff --git a/anonstream/routes/wrappers.py b/anonstream/routes/wrappers.py index cd184ae..04ff2db 100644 --- a/anonstream/routes/wrappers.py +++ b/anonstream/routes/wrappers.py @@ -72,7 +72,7 @@ def auth_required(f): return wrapper def generate_and_add_user( - timestamp, token=None, broadcaster=False, verified=False, + timestamp, token=None, broadcaster=False, verified=False, headers=None, ): token = token or generate_token() user = generate_user( @@ -80,6 +80,7 @@ def generate_and_add_user( token=token, broadcaster=broadcaster, verified=verified, + headers=headers, ) USERS_BY_TOKEN[token] = user USERS_UPDATE_BUFFER.add(token) @@ -132,6 +133,7 @@ def with_user_from(context, fallback_to_token=False): if CONFIG['ACCESS_CAPTCHA'] and not broadcaster: if user is not None: user['last']['seen'] = timestamp + user['headers'] = tuple(context.headers) response = await f(timestamp, user, *args, **kwargs) elif fallback_to_token: #assert not broadcaster @@ -146,8 +148,14 @@ def with_user_from(context, fallback_to_token=False): else: if user is not None: user['last']['seen'] = timestamp + user['headers'] = tuple(context.headers) else: - user = generate_and_add_user(timestamp, token, broadcaster) + user = generate_and_add_user( + timestamp, + token, + broadcaster, + headers=tuple(context.headers), + ) response = await f(timestamp, user, *args, **kwargs) # Set cookie