Add argparse to populate.py and fix some lints (#1064)

* Make `populate.py` configurable from command line
selenium-screenshot-testing
Kevin Chung 2019-07-21 15:08:37 -04:00 committed by GitHub
parent bf05b30d56
commit f2e0b9e8b5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 20 additions and 14 deletions

View File

@ -5,21 +5,28 @@ import datetime
import hashlib import hashlib
import random import random
import sys import sys
import argparse
from CTFd import create_app from CTFd import create_app
from CTFd.models import * from CTFd.models import Users, Teams, Challenges, Flags, Awards, ChallengeFiles, Fails, Solves
try: parser = argparse.ArgumentParser()
mode = sys.argv[1]
except IndexError: parser.add_argument("--mode", help="Set user mode", default="teams")
mode = 'teams' parser.add_argument("--users", help="Amount of users to generate", default=50, type=int)
parser.add_argument("--teams", help="Amount of teams to generate", default=10, type=int)
parser.add_argument("--challenges", help="Amount of challenges to generate", default=20, type=int)
parser.add_argument("--awards", help="Amount of awards to generate", default=5, type=int)
args = parser.parse_args()
app = create_app() app = create_app()
USER_AMOUNT = 50 mode = args.mode
TEAM_AMOUNT = 10 if mode == 'teams' else 0 USER_AMOUNT = args.users
CHAL_AMOUNT = 20 TEAM_AMOUNT = args.teams if args.mode == 'teams' else 0
AWARDS_AMOUNT = 5 CHAL_AMOUNT = args.challenges
AWARDS_AMOUNT = args.awards
categories = [ categories = [
'Exploitation', 'Exploitation',
@ -201,7 +208,7 @@ def gen_name():
def gen_team_name(): def gen_team_name():
return random.choice(hipsters).capitalize() + str(random.randint(1,1000)) return random.choice(hipsters).capitalize() + str(random.randint(1, 1000))
def gen_email(): def gen_email():
@ -302,7 +309,6 @@ if __name__ == '__main__':
db.session.commit() db.session.commit()
# Generating Users # Generating Users
print("GENERATING USERS") print("GENERATING USERS")
used = [] used = []
@ -331,7 +337,7 @@ if __name__ == '__main__':
user.team_id = random.randint(1, TEAM_AMOUNT) user.team_id = random.randint(1, TEAM_AMOUNT)
db.session.add(user) db.session.add(user)
count += 1 count += 1
except: except Exception:
pass pass
db.session.commit() db.session.commit()
@ -356,7 +362,7 @@ if __name__ == '__main__':
chalid = random.randint(1, CHAL_AMOUNT) chalid = random.randint(1, CHAL_AMOUNT)
if chalid not in used: if chalid not in used:
used.append(chalid) used.append(chalid)
user = Users.query.filter_by(id=x+1).first() user = Users.query.filter_by(id=x + 1).first()
solve = Solves( solve = Solves(
user_id=user.id, user_id=user.id,
team_id=user.team_id, team_id=user.team_id,
@ -432,7 +438,7 @@ if __name__ == '__main__':
chalid = random.randint(1, CHAL_AMOUNT) chalid = random.randint(1, CHAL_AMOUNT)
if chalid not in used: if chalid not in used:
used.append(chalid) used.append(chalid)
user = Users.query.filter_by(id=x+1).first() user = Users.query.filter_by(id=x + 1).first()
wrong = Fails( wrong = Fails(
user_id=user.id, user_id=user.id,
team_id=user.team_id, team_id=user.team_id,