95 lines
2.6 KiB
Python
Executable File
95 lines
2.6 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""CLI tool to manage users in the KOReaderServerFetcher database."""
|
|
|
|
import argparse
|
|
import sqlite3
|
|
import secrets
|
|
import sys
|
|
|
|
|
|
def add_user(db_path, username):
|
|
"""Add a new user with a random token."""
|
|
token = secrets.token_urlsafe(32)
|
|
|
|
conn = sqlite3.connect(db_path)
|
|
cursor = conn.cursor()
|
|
|
|
# Ensure table exists
|
|
cursor.execute('''
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
auth_token TEXT NOT NULL UNIQUE,
|
|
name TEXT UNIQUE,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
''')
|
|
|
|
try:
|
|
cursor.execute(
|
|
'INSERT INTO users (auth_token, name) VALUES (?, ?)',
|
|
(token, username)
|
|
)
|
|
conn.commit()
|
|
print(f"Created user: {username}")
|
|
print(f"Token: {token}")
|
|
return True
|
|
except sqlite3.IntegrityError:
|
|
print(f"Error: Username '{username}' already exists", file=sys.stderr)
|
|
return False
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def delete_user(db_path, username):
|
|
"""Delete a user by their username."""
|
|
conn = sqlite3.connect(db_path)
|
|
cursor = conn.cursor()
|
|
|
|
# First check if user exists
|
|
cursor.execute('SELECT id, auth_token FROM users WHERE name = ?', (username,))
|
|
result = cursor.fetchone()
|
|
|
|
if not result:
|
|
print(f"Error: No user found with username '{username}'", file=sys.stderr)
|
|
conn.close()
|
|
return False
|
|
|
|
user_id, token = result
|
|
|
|
# Delete user's downloads first
|
|
cursor.execute('DELETE FROM downloads WHERE user_id = ?', (user_id,))
|
|
downloads_deleted = cursor.rowcount
|
|
|
|
# Delete the user
|
|
cursor.execute('DELETE FROM users WHERE id = ?', (user_id,))
|
|
conn.commit()
|
|
conn.close()
|
|
|
|
print(f"Deleted user: {username}")
|
|
if downloads_deleted > 0:
|
|
print(f"Also removed {downloads_deleted} download record(s)")
|
|
return True
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Manage users in the KOReaderServerFetcher database')
|
|
parser.add_argument('--db', required=True, help='Path to the SQLite database')
|
|
parser.add_argument('--username', help='Username for the user')
|
|
parser.add_argument('--delete', action='store_true', help='Delete a user instead of adding')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not args.username:
|
|
parser.error('--username is required')
|
|
|
|
if args.delete:
|
|
success = delete_user(args.db, args.username)
|
|
else:
|
|
success = add_user(args.db, args.username)
|
|
|
|
sys.exit(0 if success else 1)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|