generate_cu_data.py

#!/usr/bin/env python
'''Generate SQL statements for CU'''


from faker import Faker
import bcrypt
import random
import json
import os
import sys
import tempfile

TIMEFORMAT = '%Y-%m-%d %H:%M:%S '


def generate_test_data(N=1000):
    '''Generate test data'''
    CU = 'gatlingcu'
    FAKER = Faker()
    (gatling_data, sql) = generate_profiles(FAKER, CU, N)
    tmpdir = tempfile.gettempdir()
    with open(os.path.join(tmpdir, "gatling.users.json"), "w") as users_json:
        print(json.dumps(gatling_data, indent=4), file=users_json)
    with open(os.path.join(tmpdir, "gatling.users.sql"), "w") as users_sql:
        for s in sql:
            print(s, file=users_sql)
            print('\n\n\n', file=users_sql)


def generate_profiles(faker, cu, N=1000):
    '''Generate profiles'''
    USERNAMES = {}

    def generate():
        '''do generation'''
        for i in range(0, N):
            profile = generate_cuuser(i, faker, cu, USERNAMES)
            yield (gatling_test_data(profile),
                   generate_profile_sql(profile))

    profiles = list(generate())
    return ([x[0] for x in profiles], [x[1] for x in profiles])


def generate_profile_sql(profile):
    '''Generate CU SQL inserts'''
    sqls = [
        generate_cuuser_sql(profile),
        generate_cugroup_sql(profile),
        generate_cumemberacct_sql(profile),
        generate_cuaccountbalance_sql(profile),
        generate_cuaccounthistory_sql(profile),
        generate_cumemberacctrights_sql(profile),
        generate_cuuseraccounts_sql(profile),
    ]
    return '\n'.join(sqls)


def bool_to_str(b):
    '''return Y or N, ugh...'''
    return 'Y' if b else 'N'


def random_bool_to_str(faker):
    '''return random bool as str'''
    return bool_to_str(faker.boolean())


def generate_cuuser(index, faker, cu, USERNAMES):
    '''Generate CU user information, return dictionary'''
    r = random.Random()
    # Ensure we are generating unique profiles
    while True:
        profile = faker.simple_profile(sex=None)
        if USERNAMES.get(profile['username']) is None:
            USERNAMES[profile['username']] = 1
            break
    uid = index + 1
    account_balance = r.random() * r.randint(10, 100)
    return {**profile, **{
        'cu': cu,
        'user_id': uid,
        'account_number': uid,
        'account_description': faker.sentence(nb_words=3),
        'account_balance': account_balance,
        'account_ytdinterest': r.random() * r.randint(1, 12),
        'account_lastyrinterest': r.random() * r.randint(1, 12),
        'account_available': account_balance - r.random() * r.randint(1, 3),
        'group': uid,
        'password': faker.password(length=10),
        'primary_accounts': uid,
        'billpayid': uid + r.randint(1, 20),
        'estmt_flag': random_bool_to_str(faker),
        'egenl_flag': 'N',
        'failedremain': 5,
        'forcechange': 'N',
        'forceremain': r.randint(0, 100),
        'lastlogin': faker.past_datetime().strftime(TIMEFORMAT),
        'priorlogin': faker.past_datetime().strftime(TIMEFORMAT),
        'pwchange': faker.past_datetime().strftime(TIMEFORMAT),
        'msg_tx': 0,
        'employee': 'N',
        'userflags': 0,
        'confidence': faker.text(max_nb_chars=20),
        'challenge_quest_id': 0,
        'mfaquest': {
            'answers': {
                '54': faker.word(),
            },
            'challenge': 0
        }
    }}


def generate_cuuser_sql(profile):
    '''Generate statement for cusuer'''
    TEMPLATE = ('INSERT INTO {cu}user VALUES('
                '{user_id},'
                '{group},'
                "'{username}',"
                "'{passwd}',"
                "'{mail}',"
                '{primary_accounts},'
                "'{egenl_flag}',"
                '{failedremain},'
                "'{forcechange}',"
                '{forceremain},'
                'NULL,'
                'NULL,'
                'NULL,'
                "'{pwchange}',"
                '{msg_tx},'
                "'{employee}',"
                '{userflags},'
                "'{confidence}',"
                '{challenge_quest_id},'
                'NULL,'
                'NULL,'
                'NULL,'
                "'{name}',"
                'NULL,'
                'NULL,'
                'NULL,'
                "'{mfaquest}');")
    return TEMPLATE.format(**{**profile, **{
        'passwd': bcrypt.hashpw(profile['password'].encode('utf8'),
                                bcrypt.gensalt()).decode('utf8'),
        'mfaquest': json.dumps(profile['mfaquest']),
    }})


def generate_cugroup_sql(profile):
    '''Generate statement for cugroup'''
    TEMPLATE = ("INSERT INTO {cu}group VALUES ({group}, '{group}', 200, 0);")
    return TEMPLATE.format(**profile)


def generate_cumemberacct_sql(profile):
    '''Generate statement for cu member accounts'''
    TEMPLATE = ('INSERT INTO {cu}memberacct VALUES('
                '{user_id},'
                "'{account_number}',"
                "'{estmt_flag}',"
                '{billpayid},'
                '0,'
                "'L',"
                'NULL,'
                'NULL);')
    return TEMPLATE.format(**profile)


def generate_cuaccountbalance_sql(profile):
    '''Generate INSERT statement for cu member account balance'''
    TEMPLATE = ('INSERT INTO {cu}accountbalance VALUES('
                "'{account_number}',"
                "'502',"
                '0,'
                "'Y',"
                "'{account_description}',"
                '{account_balance},'
                '{account_ytdinterest},'
                '{account_lastyrinterest},'
                '{account_available},'
                'NULL,'
                'true,'
                'true,'
                'NULL,'
                'NULL,'
                "'',"
                'NULL,'
                'NULL,'
                'NULL,'
                'NULL,'
                'NULL);')
    return TEMPLATE.format(**profile)


def generate_cuaccounthistory_sql(profile):
    '''Generate SQL insert statements for account history'''
    FAKER = Faker()
    r = random.Random()
    TEMPLATE = ('INSERT INTO {cu}accounthistory VALUES('
                "'{account_number}',"
                "'502',"
                '0,'
                "'{trace_number}',"
                "'0',"
                "'{date}',"
                '{amount},'
                "'{description}',"
                '{balance},'
                'NULL,'
                "'{sortkey}');")

    def generate():
        '''do generation'''
        INT_MAX = int((2 ** 32 - 1) / 2)
        for i in range(1, r.randint(2, 6)):
            amount = r.random() * r.randint(10, 20)
            yield TEMPLATE.format(**{**profile, **{
                'trace_number': str(r.randint(1, INT_MAX)).zfill(20),
                'date': FAKER.past_datetime().strftime(TIMEFORMAT),
                'amount': amount,
                'balance': profile['account_balance'] - amount,
                'description': FAKER.sentence(nb_words=6),
                'sortkey': str(r.randint(1, INT_MAX)).zfill(16),
            }})
    return '\n'.join(generate())


def generate_cumemberacctrights_sql(profile):
    '''Generate statement for cu member account rights'''
    TEMPLATE = ('INSERT INTO {cu}memberacctrights VALUES('
                '{user_id},'
                "'{account_number}',"
                "'ACCESS',"
                'true,'
                'NULL);')
    return TEMPLATE.format(**profile)


def generate_cuuseraccounts_sql(profile):
    '''Generate statement for cu user accounts'''
    TEMPLATE = ('INSERT INTO {cu}useraccounts VALUES('
                '{user_id},'
                "'{account_number}',"
                "'502',"
                '0,'
                "'D',"
                "'',"
                'true,'
                'true,'
                'true,'
                'true,'
                'true,'
                'true,'
                '0,'
                '30,'
                "'D');")
    return TEMPLATE.format(**profile)


def gatling_test_data(profile):
    '''filter fields for gatling test harness'''
    KEYS = {
        'username': None,
        'mail': lambda x: ('email', x),
        'password': None,
        'cu': lambda x: ('cu', x.upper()),
        'account_number': lambda x: ('account_id', x),
        'mfaquest': lambda x: ('mfa', x['answers']['54']),
    }

    def filter_keys(profile, keys):
        '''Filter out keys, based on KEYS dictionary'''
        for k, f in keys.items():
            if not f:
                yield (k, profile[k])
            else:
                yield f(profile[k])
    return dict(filter_keys(profile, KEYS))


if __name__ == '__main__':
    N = 1000
    if len(sys.argv[1:]) > 0:
        N = int(sys.argv[1])
    generate_test_data(N)