LCB/database.py

import sqlite3
from settings import SESSIONS_FOLDER, DATABASE_FILE
from telethon import TelegramClient
from settings import SEP
from Console import console
from time import time
import os


def create_new_main_db():
    try:
        db = sqlite3.connect(DATABASE_FILE)
    except sqlite3.OperationalError:
        os.mkdir("data")
        db = sqlite3.connect(DATABASE_FILE)

    cur = db.cursor()
    # Создаем таблицу
    cur.execute("""CREATE TABLE IF NOT EXISTS Account(
        ID INTEGER PRIMARY KEY,
        PHONE TEXT,
        PASS TEXT,
        API_ID TEXT,
        API_HASH TEXT,
        IS_ACTIVE INT,
        LITECOIN_WALLET TEXT,
        LITECOINPASS TEXT,
        LITECOINPHRASE TEXT, 
        SESSION TEXT,
        SAFEUMLOGIN TEXT,
        SAFEUMPASSWORD TEXT
    )""")
    cur.execute(""" CREATE TABLE  IF NOT EXISTS Proxy(
    	    ID INTEGER PRIMARY KEY,
    	    IP	TEXT,
    	    PROTOCOL TEXT,
    	    PORT TEXT,
    	    IS_ACTIVE INTEGER
            )""")
    cur.execute(""" CREATE TABLE  IF NOT EXISTS Proxy(
    	    ID INTEGER PRIMARY KEY,
    	    IP	TEXT,
    	    PROTOCOL TEXT,
    	    PORT TEXT,
    	    IS_ACTIVE INTEGER
            )""")

    cur.execute("""CREATE TABLE IF NOT EXISTS Groups(
            ID INTEGER PRIMARY KEY,
            NAME TEXT, 
            TG_LINK TEXT
            
            )""")
    cur.execute("""CREATE TABLE IF NOT EXISTS Group_Account_connection(
                ID INTEGER PRIMARY KEY,
                ACCOUNT_ID INTEGER NOT NULL, 
                GROUP_ID INTEGER NOT NULL,
                TIME_OF_LEAVING TEXT,
                FOREIGN KEY(ACCOUNT_ID) REFERENCES Account(ID),
	            FOREIGN KEY(GROUP_ID) REFERENCES Groups(ID)     
                )""")

    db.commit()


def get_all_data():
    with sqlite3.connect(DATABASE_FILE) as db:
        cur = db.cursor()
        cur.execute("SELECT SESSION, API_ID, API_HASH, IS_ACTIVE, LITECOIN_WALLET, ID, PHONE FROM Account")

        return cur.fetchall()


def get_proxys():
    with sqlite3.connect(DATABASE_FILE) as db:
        cur = db.cursor()
        cur.execute("""SELECT IP, PROTOCOL, PORT, IS_ACTIVE FROM Proxy""")

        ret = []
        for el in cur.fetchall():
            ret.append(el)

    return ret


def add_new_colums():
    with sqlite3.connect(DATABASE_FILE) as db:
        cur = db.cursor()
        cur.execute("""
        ALTER TABLE Account ADD COLUMN SESSION TEXT        
        """)
        db.commit()


def add_account(phone, password, api_id, api_hash, is_working, wallet, wallet_pass,
                recover_phrase, safelogin, password_2):
    db = sqlite3.connect(DATABASE_FILE)
    cur = db.cursor()

    ret = cur.execute("SELECT ID FROM Account").fetchall()
    if ret:
        new_id_account = ret[-1][0] + 1
    else:
        new_id_account = 1

    if api_id and api_hash:
        session = SESSIONS_FOLDER + SEP + "anon" + str(new_id_account) + '.session'
        client = TelegramClient(session, int(api_id), api_hash)
        client.start()
    else:
        session = ''

    cur.execute("""INSERT INTO Account(ID, PHONE, PASS, API_ID, API_HASH, IS_ACTIVE, LITECOIN_WALLET, SESSION,
        LITECOINPASS, LITECOINPHRASE, SAFEUMLOGIN, SAFEUMPASSWORD)
        VALUES (?,?,?,?,?,?,?,?,?,?, ?, ?);""",
                (new_id_account, phone, password, api_id, api_hash, is_working, wallet, session, wallet_pass,
                 recover_phrase, safelogin, password_2))
    db.commit()
    db.close()


def add_group(name, tg_link):
    db = sqlite3.connect(DATABASE_FILE)
    cur = db.cursor()

    ret = cur.execute("SELECT ID FROM Groups WHERE TG_LINK = ?", (tg_link,)).fetchall()
    if ret:
        return  # group already exist in database

    ret = cur.execute("SELECT ID FROM Groups").fetchall()
    if ret:
        new_id = ret[-1][0] + 1
    else:
        new_id = 1

    cur.execute("""
    INSERT INTO Groups(ID, NAME, TG_LINK) VALUES (?, ?, ?)
    """, (new_id, name, tg_link))

    db.commit()
    db.close()


def add_group_account_connection(client, group_link, stay_in_group=1):
    db = sqlite3.connect(DATABASE_FILE)
    cur = db.cursor()

    ret = cur.execute("SELECT ID FROM Group_Account_connection").fetchall()
    if ret:
        new_id = ret[-1][0] + 1
    else:
        new_id = 1
    client_id = cur.execute("SELECT ID FROM Account WHERE  API_ID = ?", (client.api_id,)).fetchall()
    if client_id:
        client_id = client_id[0]
    else:
        console.send_alert("Error when executing SELECT ID FROM Account WHERE  API_ID = {}")
        return

    group_id = cur.execute("SELECT ID FROM Groups WHERE TG_LINK = ?", (group_link,)).fetchall()
    if group_id:
        group_id = group_id[0]
    else:
        console.send_alert("Error when executing SELECT ID FROM Groups WHERE TG_LINK = {}")
        return

    cur.execute("INSERT INTO Group_Account_connection(ID, ACCOUNT_ID, GROUP_ID, TIME_OF_LEAVING) VALUES (?, ?, ?, ?)",
                (new_id, client_id[0], group_id[0], str(int(time()) + stay_in_group * 3600)))

    db.commit()


def delete_old_connections(client):
    db = sqlite3.connect(DATABASE_FILE)
    cur = db.cursor()

    client_id = cur.execute("SELECT ID FROM Account WHERE  API_ID = {}".format(client.api_id)).fetchall()
    if client_id:
        client_id = client_id[0]
    else:
        console.send_alert("Error when executing SELECT ID FROM Account WHERE  API_ID = {}")
        return

    ret = cur.execute("SELECT ID, TIME_OF_LEAVING, GROUP_ID FROM Group_Account_connection WHERE ACCOUNT_ID = ?",
                      (client_id[0],)).fetchall()
    groups_to_leave = []
    for el in ret:

        if time() > int(el[1]):
            cur.execute("DELETE FROM Group_Account_connection WHERE ID = ?", (el[0],))
            groups_to_leave.append(el[2])

    tg_links_to_leave = []
    for el in groups_to_leave:
        ret = cur.execute("SELECT NAME FROM Groups WHERE ID = ?", (el,)).fetchall()

        tg_links_to_leave.append(ret[0][0])
    db.commit()
    return tg_links_to_leave


if __name__ == '__main__':
    create_new_main_db()