diff --git a/zerver/management/commands/rate_limit.py b/zerver/management/commands/rate_limit.py index 66954ef257..b132c8507a 100644 --- a/zerver/management/commands/rate_limit.py +++ b/zerver/management/commands/rate_limit.py @@ -4,13 +4,13 @@ from __future__ import print_function from typing import Any from argparse import ArgumentParser -from zerver.models import UserProfile, get_user_profile_by_email +from zerver.models import UserProfile from zerver.lib.rate_limiter import block_user, unblock_user +from zerver.lib.management import ZulipBaseCommand -from django.core.management.base import BaseCommand from optparse import make_option -class Command(BaseCommand): +class Command(ZulipBaseCommand): help = """Manually block or unblock a user from accessing the API""" def add_arguments(self, parser): @@ -37,6 +37,7 @@ class Command(BaseCommand): help="Whether or not to also block all bots for this user.") parser.add_argument('operation', metavar='', type=str, choices=['block', 'unblock'], help="operation to perform (block or unblock)") + self.add_realm_args(parser) def handle(self, *args, **options): # type: (*Any, **Any) -> None @@ -45,8 +46,9 @@ class Command(BaseCommand): print("Please enter either an email or API key to manage") exit(1) + realm = self.get_realm(options) if options['email']: - user_profile = get_user_profile_by_email(options['email']) + user_profile = self.get_user(options['email'], realm) else: try: user_profile = UserProfile.objects.get(api_key=options['api_key'])