models.py 5.7 KB
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os

from sqlalchemy import Column, Integer, String, create_engine, func
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import SingletonThreadPool

from conf.settings import user_db_path
from lib import common

Base = declarative_base()
metadata = Base.metadata

engine = create_engine('sqlite:///{}'.format(os.path.join(user_db_path, "userlist.db")),
                       poolclass=SingletonThreadPool,
                       connect_args={'check_same_thread': False})


def NewSession():
    DbSession = sessionmaker(bind=engine)
    session = DbSession()
    return session


session = NewSession()


class Config(Base):
    __tablename__ = 'config'
    id = Column(Integer, primary_key=True, autoincrement=True)
    key = Column(String(50), index=True, nullable=False, unique=True)
    value = Column(String(256))

    @staticmethod
    def get(key, value_func=None):
        conf = session.query(Config).filter_by(key=key).first()
        if not conf:
            if value_func:
                conf = Config(key=key, value=value_func())
                try:
                    session.add(conf)
                    session.commit()
                except:
                    session.rollback()

                return conf.value
            return None
        return conf.value

    @staticmethod
    def set(key, value):
        conf = session.query(Config).filter_by(key=key).first()
        if not conf:
            conf = Config(key=key, value=value)
            try:
                session.add(conf)
                session.commit()
            except:
                session.rollback()

            return conf
        else:
            try:
                conf.value = value
                session.commit()
            except:
                session.rollback()
        return conf

    def __repr__(self):
        return "Config(key={}, value={})".format(self.key, self.value)


class UserList(Base):
    __tablename__ = 'users'

    id = Column(Integer, primary_key=True, autoincrement=True)
    email = Column(String(50), index=True, nullable=False, unique=True)
    password = Column(String(50), nullable=False)
    cookie = Column(String(512))
    user_agent = Column(String(256))
    fbid = Column(String(20), index=True)
    status = Column(Integer, default=0, nullable=False, index=True)
    proxy = Column(String(256))
    token = Column(String(256))
    phone_user_agent = Column(String(256))
    phone_cookie = Column(String(512))
    phone_token = Column(String(256))
    adId = Column(String(256))
    deviceId = Column(String(256))
    familyDeviceId = Column(String(256))
    machineId = Column(String(256))

    def __repr__(self):
        if self.proxy:
            proxies = common.frombase64(self.proxy)
            proxies = f'{proxies.get("host")}:{proxies.get("port")}'
        else:
            proxies = None

        return "User(id={}, email={}, password={}, cookie={}, fbid={}, status={}, proxy={})" \
            .format(self.id, self.email, self.password, len(self.cookie) if self.cookie else None, self.fbid,
                    self.status, proxies)

    def format_cookie(self):
        if self.cookie:
            return dict([tuple(sub.split("=")) for sub in self.cookie.split('; ') if sub])
        else:
            return {}

    def set(self, **kwargs):
        for k, v in kwargs.items():
            if isinstance(v, (list, dict)):
                continue
            setattr(self, k, v)
        try:
            session.commit()
        except:
            session.rollback()

        return self

    @staticmethod
    def update(**kwargs):
        unique = tuple(kwargs)[0]
        if unique == 'email':
            user_ = session.query(UserList).filter_by(email=kwargs.get(unique)).first()
        elif unique == 'fbid':
            user_ = session.query(UserList).filter_by(fbid=kwargs.get(unique)).first()
        else:
            raise BaseException("条件不对")
        kwargs.pop(unique)
        for k, v in kwargs.items():
            setattr(user_, k, v)
        try:
            session.commit()
        except:
            session.rollback()
        return user_

    @staticmethod
    def insert(**kwargs):
        u = UserList(**kwargs)
        try:
            session.add(u)
            session.commit()
        except:
            session.rollback()
        return u

    @staticmethod
    def get(**kwargs):
        unique = tuple(kwargs)[0]
        if unique == 'email':
            user_ = session.query(UserList).filter_by(email=kwargs.get(unique)).first()
        elif unique == 'fbid':
            user_ = session.query(UserList).filter_by(fbid=kwargs.get(unique)).first()
        else:
            raise BaseException("条件不对")

        return user_

    @staticmethod
    def remove(**kwargs):
        unique = tuple(kwargs)[0]
        if unique == 'email':
            user_ = session.query(UserList).filter_by(email=kwargs.get(unique)).first()
        elif unique == 'fbid':
            user_ = session.query(UserList).filter_by(fbid=kwargs.get(unique)).first()
        else:
            raise BaseException("条件不对")
        if user_:
            try:
                session.delete(user_)
                session.commit()
            except:
                session.rollback()
        else:
            raise BaseException("用户不存在")

    @staticmethod
    def all() -> list:
        users = session.query(UserList).all()
        return users

    @staticmethod
    def query(**kwargs) -> list:
        users = session.query(UserList).filter_by(**kwargs).all()
        return users

    @staticmethod
    def all_count() -> int:
        count = session.query(func.count(UserList.id)).scalar()
        session.commit()
        return count