67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
import sqlite3
|
|
|
|
# Function to add a new user to the database
|
|
def setupMfaSecret(user, secret):
|
|
with sqlite3.connect('mfa.db') as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
INSERT INTO users (user, secret)
|
|
VALUES (?, ?)
|
|
''', (user, secret))
|
|
conn.commit()
|
|
|
|
# Function to retrieve the number corresponding to a username
|
|
def getMfaSecret(user):
|
|
with sqlite3.connect('mfa.db') as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
SELECT secret FROM users
|
|
WHERE user = ?
|
|
''', (user,))
|
|
result = cursor.fetchone()
|
|
if result:
|
|
return result[0]
|
|
else:
|
|
return None
|
|
|
|
# Function to update the number for a specific username
|
|
def updateMfaSecret(user, secret):
|
|
with sqlite3.connect('mfa.db') as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
UPDATE users
|
|
SET secret = ?
|
|
WHERE user = ?
|
|
''', (secret, user))
|
|
conn.commit()
|
|
|
|
def removeMfa(user):
|
|
with sqlite3.connect('mfa.db') as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
DELETE FROM users
|
|
WHERE user = ?
|
|
''', (user,))
|
|
conn.commit()
|
|
|
|
def initializeDb():
|
|
with sqlite3.connect('mfa.db') as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
user TEXT PRIMARY KEY,
|
|
secret TEXT
|
|
)
|
|
''')
|
|
|
|
# initializeDb()
|
|
|
|
# Example usage
|
|
# add_user('user1', 123)
|
|
# add_user('user2', 456)
|
|
|
|
# print(f"Number for 'user1': {get_number('user1')}")
|
|
|
|
# update_number('user1', 789)
|
|
# print(f"Updated number for 'user1': {get_number('user1')}")
|