140 lines
5.4 KiB
Python
140 lines
5.4 KiB
Python
# server/db_setup.py
|
|
import sqlite3
|
|
import hashlib
|
|
import os
|
|
from contextlib import closing
|
|
|
|
DATABASE = 'transfer_service.db'
|
|
|
|
def initialize_db():
|
|
with closing(sqlite3.connect(DATABASE)) as conn, conn, closing(conn.cursor()) as c:
|
|
# Create users table with salt
|
|
c.execute('''
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id INTEGER PRIMARY KEY,
|
|
username TEXT UNIQUE NOT NULL,
|
|
password TEXT NOT NULL,
|
|
salt TEXT NOT NULL,
|
|
login_attempts INTEGER DEFAULT 0
|
|
)
|
|
''')
|
|
|
|
# Create uploads table for storing links, files, and images
|
|
c.execute('''
|
|
CREATE TABLE IF NOT EXISTS uploads (
|
|
id INTEGER PRIMARY KEY,
|
|
uploader TEXT NOT NULL,
|
|
file_type TEXT NOT NULL, -- 'link', 'file', or 'image'
|
|
content TEXT NOT NULL, -- The actual link, filename, or file path
|
|
uploaded_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
# Check if the 'salt' column exists in the users table and add it if missing
|
|
c.execute("PRAGMA table_info(users)")
|
|
columns = [column[1] for column in c.fetchall()]
|
|
if 'salt' not in columns:
|
|
c.execute("ALTER TABLE users ADD COLUMN salt TEXT")
|
|
update_existing_users_with_salts() # Update existing users with salts
|
|
|
|
conn.commit()
|
|
|
|
def generate_salt():
|
|
return os.urandom(16).hex()
|
|
|
|
def hash_password(password, salt):
|
|
return hashlib.sha256((password + salt).encode()).hexdigest()
|
|
|
|
def update_existing_users_with_salts():
|
|
"""
|
|
Updates existing users to include a unique salt and rehashes their passwords.
|
|
"""
|
|
with closing(sqlite3.connect(DATABASE)) as conn, closing(conn.cursor()) as c:
|
|
c.execute('SELECT id, password FROM users')
|
|
users = c.fetchall()
|
|
|
|
for user_id, password in users:
|
|
salt = generate_salt()
|
|
hashed_password = hash_password(password, salt)
|
|
c.execute('UPDATE users SET password = ?, salt = ? WHERE id = ?', (hashed_password, salt, user_id))
|
|
|
|
conn.commit()
|
|
print("Updated existing users with salts.")
|
|
|
|
def add_user(username, password):
|
|
# Generate a unique salt for each user
|
|
salt = os.urandom(16)
|
|
hashed_password = hashlib.sha256(salt + password.encode()).hexdigest()
|
|
|
|
try:
|
|
with closing(sqlite3.connect(DATABASE)) as conn, conn, closing(conn.cursor()) as c:
|
|
c.execute('INSERT INTO users (username, password, salt) VALUES (?, ?, ?)', (username, hashed_password, salt))
|
|
conn.commit()
|
|
print(f"User '{username}' added successfully.")
|
|
except sqlite3.IntegrityError:
|
|
print(f"User '{username}' already exists.")
|
|
|
|
def delete_user(username):
|
|
"""
|
|
Deletes a user from the database based on the provided username.
|
|
"""
|
|
try:
|
|
with closing(sqlite3.connect(DATABASE)) as conn, conn, closing(conn.cursor()) as c:
|
|
c.execute('DELETE FROM users WHERE username = ?', (username,))
|
|
conn.commit()
|
|
print(f"User '{username}' deleted successfully.")
|
|
except sqlite3.Error as e:
|
|
print(f"Error deleting user '{username}': {e}")
|
|
|
|
def get_user(username):
|
|
with closing(sqlite3.connect(DATABASE)) as conn, closing(conn.cursor()) as c:
|
|
# Select only password, salt, and login_attempts
|
|
c.execute('SELECT password, salt, login_attempts FROM users WHERE username = ?', (username,))
|
|
return c.fetchone()
|
|
|
|
def reset_login_attempts(username):
|
|
with closing(sqlite3.connect(DATABASE)) as conn, conn, closing(conn.cursor()) as c:
|
|
c.execute('UPDATE users SET login_attempts = 0 WHERE username = ?', (username,))
|
|
conn.commit()
|
|
|
|
def increment_login_attempts(username):
|
|
with closing(sqlite3.connect(DATABASE)) as conn, conn, closing(conn.cursor()) as c:
|
|
if username:
|
|
c.execute('UPDATE users SET login_attempts = login_attempts + 1 WHERE username = ?', (username,))
|
|
c.execute('SELECT login_attempts FROM users WHERE username = ?', (username,))
|
|
login_attempts = c.fetchone()[0]
|
|
conn.commit()
|
|
return login_attempts
|
|
else:
|
|
return None
|
|
|
|
def add_upload(uploader, file_type, content):
|
|
with closing(sqlite3.connect(DATABASE)) as conn, conn, closing(conn.cursor()) as c:
|
|
c.execute('INSERT INTO uploads (uploader, file_type, content) VALUES (?, ?, ?)', (uploader, file_type, content))
|
|
conn.commit()
|
|
|
|
def get_uploads():
|
|
with closing(sqlite3.connect(DATABASE)) as conn, closing(conn.cursor()) as c:
|
|
c.execute('SELECT id, uploader, file_type, content FROM uploads')
|
|
return c.fetchall()
|
|
|
|
def delete_upload(upload_id):
|
|
with closing(sqlite3.connect(DATABASE)) as conn, conn, closing(conn.cursor()) as c:
|
|
c.execute('DELETE FROM uploads WHERE id = ?', (upload_id,))
|
|
conn.commit()
|
|
|
|
def update_upload_filename(upload_id, new_name):
|
|
with closing(sqlite3.connect(DATABASE)) as conn, closing(conn.cursor()) as c:
|
|
c.execute('UPDATE uploads SET content = ? WHERE id = ?', (new_name, upload_id))
|
|
conn.commit()
|
|
|
|
if __name__ == '__main__':
|
|
initialize_db()
|
|
|
|
# Example of initializing users (only run manually)
|
|
# add_user('iphone_user', 'your_secure_password')
|
|
# add_user('laptop_user', 'your_secure_password')
|
|
# Example of deleting user
|
|
# delete_user('iphone_user')
|
|
# delete_user('laptop_user')
|