Loadtesting with Firefox Account

get_random_token.py

#!/usr/bin/env python
from __future__ import print_function
import base64
import hmac
import os
import sys
from six.moves.urllib.parse import urlparse

from fxa import core
from fxa import errors
from fxa import oauth
from fxa.plugins.requests import DEFAULT_CLIENT_ID  # Firefox dev Client ID

FXA_API_URL = "https://api-accounts.stage.mozaws.net/v1"
FXA_OAUTH_URL = "https://oauth.stage.mozaws.net/v1"


# Constants
FXA_ERROR_ACCOUNT_EXISTS = 101


def get_bearer_token(email, password, scopes=None,
                     account_server_url=FXA_API_URL,
                     oauth_server_url=FXA_OAUTH_URL,
                     client_id=DEFAULT_CLIENT_ID):

    if scopes is None:
        scopes = ['profile']

    client = core.Client(server_url=account_server_url)
    session = client.login(email, password)

    url = urlparse(oauth_server_url)
    audience = "%s://%s/" % (url.scheme, url.netloc)

    bid_assertion = session.get_identity_assertion(audience)
    oauth_client = oauth.Client(server_url=oauth_server_url)
    token = oauth_client.authorize_token(bid_assertion,
                                         ' '.join(scopes),
                                         client_id)
    return token


def create_new_fxa_account(fxa_user_salt=None, account_server_url=FXA_API_URL):
    if 'stage' in account_server_url:
        if not fxa_user_salt:
            fxa_user_salt = os.urandom(36)
        else:
            fxa_user_salt = base64.urlsafe_b64decode(fxa_user_salt)

        password = hmac.new(fxa_user_salt, b"loadtest").hexdigest()
        email = "syncto-%s@restmail.net" % password

        client = core.Client(server_url=account_server_url)

        try:
            client.create_account(email,
                                  password=password,
                                  preVerified=True)
        except errors.ClientError as e:
            if e.errno != FXA_ERROR_ACCOUNT_EXISTS:
                raise
        finally:
            return email, password
    else:
        message = ("You are not using stage (%s), make sure your FxA test "
                   "account exists: https://123done-prod.dev.lcip.org/")
        print(message % account_server_url, file=sys.stderr)
        sys.exit(1)

if __name__ == '__main__':
    if len(sys.argv) < 1:
        msg = ("USAGE: %s [FXA_USER_SALT SCOPES "
               "ACCOUNT_SERVER_URL OAUTH_SERVER_URL CLIENT_ID]")
        print(msg % sys.argv[0], file=sys.stderr)
        sys.exit(1)

    if len(sys.argv) >= 2:
        fxa_user_salt = sys.argv[1].strip()
    else:
        fxa_user_salt = None

    if len(sys.argv) >= 3:
        scopes = [s.strip() for s in sys.argv[2].split()]
    else:
        scopes = ['profile']

    if len(sys.argv) >= 4:
        account_server_url = sys.argv[3]
    else:
        account_server_url = FXA_API_URL

    if len(sys.argv) >= 5:
        oauth_server_url = sys.argv[4]
    else:
        oauth_server_url = FXA_OAUTH_URL

    if len(sys.argv) >= 6:
        client_id = sys.argv[5]
    else:
        client_id = DEFAULT_CLIENT_ID

    print('# Creating the account...', end='', file=sys.stderr)
    sys.stderr.flush()
    email, password = create_new_fxa_account(fxa_user_salt, account_server_url)
    print("\b\b\b\t [OK]", file=sys.stderr)
    print('# Creating the token...', end='', file=sys.stderr)
    sys.stderr.flush()
    token = get_bearer_token(email, password, scopes, account_server_url,
                             oauth_server_url, client_id)
    print("\b\b\b\t [OK]", file=sys.stderr)

    print('\r# ---- BEARER TOKEN INFO ----')
    print('# User: %s' % email)
    print('# Scopes: %s' % ' '.join(scopes))
    print('# Account: %s' % account_server_url)
    print('# Oauth: %s' % oauth_server_url)
    print('# Client ID: %s' % client_id)
    print('# ---------------------------')
    print('export OAUTH_BEARER_TOKEN="%s"' % token)