initdb.py
4.13 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
#!/usr/bin/env python3
# base
import csv
import argparse
import re
import string
from sys import exit
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import cpu_count
import asyncio
# installed packages
import bcrypt
import sqlalchemy as sa
# this project
from models import Base, Student
pool = ThreadPoolExecutor() #cpu_count()
# replace password by hash for a single student dict
def hashpw(student):
pw = student.get('pw', student['uid']).encode('utf-8')
print('.', end='', flush=True)
hashed_pw = bcrypt.hashpw(pw, bcrypt.gensalt())
student['pw'] = hashed_pw
async def hash_all_passwords(executor, students):
loop = asyncio.get_event_loop()
tasks = [loop.run_in_executor(executor, hashpw, s) for s in students]
await asyncio.wait(tasks) # block until all tasks are done
print()
# SIIUE names have alien strings like "(TE)" and are sometimes capitalized
# We remove them so that students dont keep asking what it means
def fix(name):
return string.capwords(re.sub('\(.*\)', '', name).strip())
# ===========================================================================
# Parse command line options
def parse_commandline_arguments():
argparser = argparse.ArgumentParser(
description='Create new database from a CSV file (SIIUE format)')
argparser.add_argument('--db',
default='students.db',
type=str,
help='database filename')
argparser.add_argument('--demo',
action='store_true',
help='initialize database with a few fake students')
# FIXME
# argparser.add_argument('--pw',
# default='',
# type=str,
# help='default password')
argparser.add_argument('csvfile',
nargs='?',
type=str,
default='',
help='CSV filename')
return argparser.parse_args()
# ===========================================================================
def get_students_from_csv(filename):
try:
csvreader = csv.DictReader(open(filename, encoding='iso-8859-1'), delimiter=';', quotechar='"', skipinitialspace=True)
except EnvironmentError:
print(f'Error: CSV file "{filename}" not found.')
exit(1)
students = [{
'uid': s['N.º'],
'name': fix(s['Nome'])
} for s in csvreader]
return students
# ===========================================================================
args = parse_commandline_arguments()
if args.csvfile:
students = get_students_from_csv(args.csvfile)
elif args.demo:
students = [
{'uid': '1915', 'name': 'Alan Turing'},
{'uid': '1938', 'name': 'Donald Knuth'},
{'uid': '1815', 'name': 'Ada Lovelace'},
{'uid': '1969', 'name': 'Linus Torvalds'},
{'uid': '1955', 'name': 'Tim Burners-Lee'},
{'uid': '1916', 'name': 'Claude Shannon'},
{'uid': '1903', 'name': 'John von Neumann'}]
students.append({'uid': '0', 'name': 'Admin'})
print(f'Generating {len(students)} bcrypt password hashes.')
executor = ThreadPoolExecutor(cpu_count())
event_loop = asyncio.get_event_loop()
event_loop.run_until_complete(hash_all_passwords(executor, students))
event_loop.close()
print(f'Creating database: {args.db}')
engine = sa.create_engine(f'sqlite:///{args.db}', echo=False)
Base.metadata.create_all(engine) # Criate schema if needed
Session = sa.orm.sessionmaker(bind=engine)
try:
# --- start db session ---
session = Session()
session.add_all([Student(id=s['uid'], name=s['name'], password=s['pw'])
for s in students])
n = session.query(Student).count()
print(f'{n} user(s):')
users = session.query(Student).order_by(Student.id).all()
print(f' {users[0].id:8} - {users[0].name} (administrator)')
if n > 1:
print(f' {users[1].id:8} - {users[1].name}')
if n > 3:
print(' ... ...')
if n > 2:
print(f' {users[-1].id:8} - {users[-1].name}')
except sa.exc.IntegrityError:
print('!!! Integrity error !!!')
session.rollback()
except Exception as e:
print(f'Error: Database "{args.db}" already exists?')
session.rollback()
raise e
# exit(1)
else:
# --- end session ---
session.commit()