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"" 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"" 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")