rate_limit: Replace get_user_profile_by_email with get_user.

This commit is contained in:
Vishnu Ks
2017-07-08 21:28:54 +05:30
committed by Tim Abbott
parent 1ce6b8db8c
commit afc41dfb6b

View File

@@ -4,13 +4,13 @@ from __future__ import print_function
from typing import Any from typing import Any
from argparse import ArgumentParser 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.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 from optparse import make_option
class Command(BaseCommand): class Command(ZulipBaseCommand):
help = """Manually block or unblock a user from accessing the API""" help = """Manually block or unblock a user from accessing the API"""
def add_arguments(self, parser): def add_arguments(self, parser):
@@ -37,6 +37,7 @@ class Command(BaseCommand):
help="Whether or not to also block all bots for this user.") help="Whether or not to also block all bots for this user.")
parser.add_argument('operation', metavar='<operation>', type=str, choices=['block', 'unblock'], parser.add_argument('operation', metavar='<operation>', type=str, choices=['block', 'unblock'],
help="operation to perform (block or unblock)") help="operation to perform (block or unblock)")
self.add_realm_args(parser)
def handle(self, *args, **options): def handle(self, *args, **options):
# type: (*Any, **Any) -> None # type: (*Any, **Any) -> None
@@ -45,8 +46,9 @@ class Command(BaseCommand):
print("Please enter either an email or API key to manage") print("Please enter either an email or API key to manage")
exit(1) exit(1)
realm = self.get_realm(options)
if options['email']: if options['email']:
user_profile = get_user_profile_by_email(options['email']) user_profile = self.get_user(options['email'], realm)
else: else:
try: try:
user_profile = UserProfile.objects.get(api_key=options['api_key']) user_profile = UserProfile.objects.get(api_key=options['api_key'])