1
Fork 0
mirror of https://github.com/RYGhub/royalnet.git synced 2024-11-23 19:44:20 +00:00
royalnet/database.py

139 lines
No EOL
4 KiB
Python

import sqlalchemy.exc
from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, ForeignKey
from sqlalchemy.orm import sessionmaker, relationship
from sqlalchemy.ext.declarative import declarative_base
import bcrypt
class NoUsersMatchingError(Exception):
pass
class InvalidPasswordError(Exception):
pass
# Initialize the database
engine = create_engine("sqlite:///db.sqlite")
Base = declarative_base()
Session = sessionmaker(bind=engine)
class User(Base):
__tablename__ = "members"
id = Column(Integer, primary_key=True)
username = Column(String, unique=True, nullable=False)
password = Column(String, nullable=False)
telegram_id = Column(Integer, unique=True)
discord_id = Column(Integer, unique=True)
def __str__(self):
return self.username
def __repr__(self):
return f"<User {self.id}: {self.username}>"
class Diario(Base):
__tablename__ = "diario"
id = Column(Integer, primary_key=True)
text = Column(String, nullable=False)
date = Column(DateTime, nullable=False)
def __repr__(self):
return f"<Diario {self.date} {self.text}>"
Base.metadata.create_all(engine)
def create_user(username, password):
"""Create a new user and add it to the database."""
# Create a new session
session = Session()
# Hash the password with bcrypt
hashed_password = bcrypt.hashpw(password.encode("utf8"), bcrypt.gensalt())
# Create a new user
new_member = User(username=username, password=hashed_password)
# Add the newly created member to the session
session.add(new_member)
# Commit the changes
session.commit()
# TODO: check for vulnerabilities
def change_password(username, newpassword):
# Create a new session
session = Session()
# Hash the new password using bcrypt
hashed_password = bcrypt.hashpw(newpassword.encode("utf8"), bcrypt.gensalt())
# Find the user entry
users = session.query(User).filter_by(username=username).all()
if len(users) == 0:
raise NoUsersMatchingError("No users with the specified username found.")
db_user = users[0]
# Change the password and commit
db_user.password = hashed_password
session.commit()
def login(username, password, enable_exceptions=False):
"""Try to login using the database password.
The session is always returned, while the user object is returned if the login is successful."""
# Create a new session
session = Session()
# Find the matching user
db_user = session.query(User).filter_by(username=username).first()
# Verify that the user exists
if db_user is not None:
return session, None
# Test the password and return the session and the user if successful
if bcrypt.hashpw(password.encode("utf8"), db_user.password) == db_user.password:
return session, db_user
else:
if enable_exceptions:
raise InvalidPasswordError("The specified password doesn't match the user's.")
else:
return session, None
def find_user(username):
"""Find the user with the specified username and return the session and the user object."""
# Create a new session
session = Session()
# Find the matching user
db_user = session.query(User).filter_by(username=username).first()
# Return the session and the user
return session, db_user
def migrate_diario():
import datetime
session = Session()
file = open("diario.txt", encoding="utf8")
for row in file:
entry = row.split("|", 1)
new = Diario()
new.date = datetime.datetime.fromtimestamp(int(entry[0]))
new.text = entry[1]
session.add(new)
session.commit()
def new_diario_entry(dt, text):
# Create a new session
session = Session()
# Create a new diario entry
entry = Diario()
entry.date = dt
entry.text = text
# Add the entry to the database
session.add(entry)
# Commit the change
session.commit()
session = Session()
if len(session.query(User).all()) < 1:
# Look! A plaintext password!
create_user("steffo", "v3n0m-sn4k3")