mirror of https://github.com/zulip/zulip.git
74 lines
2.8 KiB
Python
74 lines
2.8 KiB
Python
from __future__ import absolute_import
|
|
|
|
from zephyr.models import UserProfile, get_user_profile_by_email
|
|
from zephyr.lib.rate_limiter import block_user, unblock_user
|
|
|
|
from django.core.management.base import BaseCommand
|
|
from django.conf import settings
|
|
from optparse import make_option
|
|
|
|
from datetime import datetime, timedelta
|
|
from collections import defaultdict
|
|
import os, time
|
|
|
|
class Command(BaseCommand):
|
|
help = """Manually block or unblock a user from accessing the API
|
|
|
|
Usage: ./manage.py rate_limit [--all-bots] [--domain all] [--seconds 60] [--api-key bf4sds] [--email f@b.com] block/unblock"""
|
|
|
|
option_list = BaseCommand.option_list + (
|
|
make_option('-e', '--email',
|
|
dest='email',
|
|
help="Email account of user."),
|
|
make_option('-a', '--api-key',
|
|
dest='api_key',
|
|
help="API key of user."),
|
|
make_option('-s', '--seconds',
|
|
dest='seconds',
|
|
default=60,
|
|
type=int,
|
|
help="Seconds to block for."),
|
|
make_option('-d', '--domain',
|
|
dest='domain',
|
|
default='all',
|
|
help="Rate-limiting domain. Defaults to 'all'."),
|
|
make_option('-b', '--all-bots',
|
|
dest='bots',
|
|
action='store_true',
|
|
default=False,
|
|
help="Whether or not to also block all bots for this user."),
|
|
)
|
|
|
|
def handle(self, *args, **options):
|
|
if len(args) == 0 or args[0] not in ('block', 'unblock'):
|
|
print "Please pass either 'block' or 'unblock"
|
|
exit(1)
|
|
|
|
if (not options['api_key'] and not options['email']) or \
|
|
(options['api_key'] and options['email']):
|
|
print "Please enter either an email or API key to manage"
|
|
exit(1)
|
|
|
|
if options['email']:
|
|
user_profile = get_user_profile_by_email(options['email'])
|
|
else:
|
|
try:
|
|
user_profile = UserProfile.objects.get(api_key=options['api_key'])
|
|
except:
|
|
print "Unable to get user profile for api key %s" % (options['api_key'], )
|
|
exit(1)
|
|
|
|
users = [user_profile]
|
|
if options['bots']:
|
|
users.extend(bot for bot in UserProfile.objects.filter(is_bot=True,
|
|
bot_owner=user_profile))
|
|
|
|
operation = args[0]
|
|
for user in users:
|
|
print "Applying operation to User ID: %s: %s" % (user.id, operation)
|
|
|
|
if operation == 'block':
|
|
block_user(user, options['seconds'], options['domain'])
|
|
elif operation == 'unblock':
|
|
unblock_user(user, options['domain'])
|