keymaster/main.py

#!/usr/bin/env python3

# Key Master v1.05
# 03/03/2024
# https://t.me/ssleg  © 2020 – 2024

import logging
import sqlite3
from datetime import datetime
from hashlib import md5, sha256
from os import path, getppid, kill
from random import randint
from signal import SIGKILL

import PySimpleGUI as Sg
import cryptocode

log_file = logging.FileHandler('key_master.log', 'w', 'utf-8')
log_file.setFormatter(logging.Formatter('%(levelname)s %(module)-13s [%(asctime)s] %(message)s'))
logging.basicConfig(level=logging.INFO, handlers=[log_file])

cc1 = 0
cc2 = 0
cc3 = 0


class LoginCard:
    """Класс хранения карточки логина/пароля"""

    __slots__ = ['__row_id', '__encrypt_key', '__hash', '__name', '__login', '__password']

    def __init__(self, encrypt_key, row_id=None, name=None, login=None, password=None):
        global cc1, cc2, cc3

        self.__row_id = row_id
        self.__encrypt_key = encrypt_key

        if row_id is None:
            self.__name = name
            self.__login = login
            self.__password = password
            control_string = self.__name + self.__login + self.__password
            self.__hash = md5(control_string.encode()).hexdigest()
            self.__write_card_to_db(mode=0)
        else:
            start = datetime.now()
            cursor.execute('select * from logins where row_id=?', (row_id,))
            row = cursor.fetchone()
            db_name = row[1]
            db_login = row[2]
            db_password = row[3]
            self.__hash = row[4]
            end = datetime.now()
            delta = end - start
            if row_id != 1:
                cc1 += delta.microseconds
                cc3 += 1

            start = datetime.now()
            self.__name = cryptocode.decrypt(db_name, encrypt_key)
            self.__login = cryptocode.decrypt(db_login, encrypt_key)
            self.__password = cryptocode.decrypt(db_password, encrypt_key)
            end = datetime.now()
            delta = end - start
            if row_id != 1:
                cc2 += delta.microseconds

    def is_valid(self):
        control_string = self.__name + self.__login + self.__password
        if isinstance(control_string, str):
            control_hash = md5(control_string.encode()).hexdigest()
            if control_hash == self.__hash:
                return True

        return False

    def __write_card_to_db(self, mode):
        db_name = cryptocode.encrypt(self.__name, self.__encrypt_key)
        db_login = cryptocode.encrypt(self.__login, self.__encrypt_key)
        db_password = cryptocode.encrypt(self.__password, self.__encrypt_key)
        if mode == 0:
            entry = (db_name, db_login, db_password, self.__hash)
            cursor.execute('insert into logins (row_name, login, password, hash) values (?,?,?,?)', entry)
            con.commit()
            cursor.execute('select row_id from logins where row_name=?', (db_name,))
            row = cursor.fetchone()
            self.__row_id = row[0]
        else:
            entry = (db_name, db_login, db_password, self.__hash, self.__row_id)
            cursor.execute('update logins set row_name=?, login=?, password=?, hash=? where row_id=?', entry)
            con.commit()

    def get_name(self):
        return self.__name

    def get_login(self):
        return self.__login

    def get_password(self):
        return self.__password

    def get_id(self):
        return self.__row_id

    def get_full_text(self):
        full_text = 'Логин:\n' + self.__login + '\n\nПароль:\n' + self.__password
        return full_text

    def update_card(self, name, login, password):
        self.__name = name
        self.__login = login
        self.__password = password
        control_string = self.__name + self.__login + self.__password
        self.__hash = md5(control_string.encode()).hexdigest()
        self.__write_card_to_db(mode=1)

    def delete_card(self):
        cursor.execute('delete from logins where row_id=?', (self.__row_id,))
        con.commit()


class AllCards:
    """Класс хранения всех карточек"""

    __slots__ = ['__encrypt_key', '__cards']

    def __init__(self, encrypt_key):
        self.__encrypt_key = encrypt_key
        row_ids = []
        cursor.execute('select row_id from logins where row_id>1')
        for row in cursor.fetchall():
            row_ids.append(row[0])

        self.__cards = []
        for row in row_ids:
            self.__cards.append(LoginCard(encrypt_key, row_id=row))
        self.sort_cards()

    def validate(self):
        error_numbers = []
        for card in self.__cards:
            valid = card.is_valid()
            if not valid:
                error_numbers.append(card.get_id())

        if len(error_numbers) > 0:
            return False, error_numbers
        else:
            return True, error_numbers

    def get_names_list(self):
        names_list = []
        for card in self.__cards:
            names_list.append(card.get_name())
        return names_list

    def add_card(self, name, login, password):
        self.__cards.append(LoginCard(self.__encrypt_key, name=name, login=login, password=password))

    def get_card_by_name(self, name):
        for card in self.__cards:
            if card.get_name() == name:
                return card.get_full_text()

    def get_card_login(self, name):
        for card in self.__cards:
            if card.get_name() == name:
                return card.get_login()

    def get_card_password(self, name):
        for card in self.__cards:
            if card.get_name() == name:
                return card.get_password()

    def get_card_id(self, name):
        for card in self.__cards:
            if card.get_name() == name:
                return card.get_id()

    @staticmethod
    def __sort_key(element: LoginCard):
        return element.get_name()

    def sort_cards(self):
        self.__cards.sort(key=self.__sort_key)

    def get_popular_logins(self):
        logins = {}
        logins_count = []
        for card in self.__cards:
            login = card.get_login()
            index = logins.get(login)
            if index is None:
                logins[login] = len(logins_count)
                logins_count.append([1, login])
            else:
                logins_count[index][0] += 1

        popular = []
        for item in logins_count:
            if item[0] > 2:
                popular.append(item[1])

        return popular

    def update_card(self, row_id, name, login, password):
        for card in self.__cards:
            if card.get_id() == row_id:
                card.update_card(name, login, password)
                break

    def delete_card(self, name):
        index = 0
        while index < len(self.__cards):
            card_name = self.__cards[index].get_name()
            if card_name == name:
                break
            index += 1
        self.__cards[index].delete_card()
        self.__cards.pop(index)


# генератор паролей
def pass_generator(species_flag=False):
    species = '!@#$%^&*()-_+=;:,./?\\|`~[]{}^' + "'" + '"<>~.'
    new_pass = ''
    species_count = 0
    for i in range(0, 16):
        cycle_flag = False
        if species_flag:
            if species_count < 4:
                probability = randint(0, 5)
                if probability == 4:
                    index = randint(0, 34)
                    new_pass += species[index]
                    species_count += 1
                    cycle_flag = True

        if not cycle_flag:
            char_index = randint(0, 35)
            if char_index > 9:
                char_index += 55
                capital = randint(0, 1)
                if capital == 0:
                    new_pass += chr(char_index)
                else:
                    new_pass += chr(char_index + 32)
            else:
                new_pass += str(char_index)

    return new_pass


# окно первого входа и задания пароля
def system_init():
    status = False
    password = ''
    layout = [[Sg.Text(
        'Задайте мастер-пароль шифрования.\nВнимание, пароль нигде не сохраняется и не может быть восстановлен!')],
        [Sg.Text('Пароль:')],
        [Sg.InputText(size=(20, 1), border_width=3)],
        [Sg.Text('Повторите пароль:')],
        [Sg.InputText(size=(20, 1), border_width=3)],
        [Sg.Button('Ok'), Sg.Text('', key='Err_text', text_color='violetred4', font='Helvetica 15 bold')]]
    window = Sg.Window("KeyMaster - приветствие", layout)
    while True:
        event, values = window.read()
        if event == Sg.WIN_CLOSED:
            break
        if event == 'Ok':
            input_word = values[0]
            input_word2 = values[1]
            if input_word == input_word2:
                if len(input_word) >= 4:
                    status = True
                    password = sha256(input_word.encode()).hexdigest()
                    break
                else:
                    err = 'Пароль слишком короткий, минимум 4 символа!'
                    window.Element('Err_text').Update(err)
            else:
                err = 'Пароли не совпадают!'
                window.Element('Err_text').Update(err)

    window.close()
    return status, password


# окно ввода и проверки пароля
def system_login():
    qt_enter_key1 = 'special 16777220'
    qt_enter_key2 = 'special 16777221'
    attempts = 3
    status = False
    password = ''
    layout = [[Sg.Text('Пароль:'), Sg.InputText(size=(20, 1), border_width=3, password_char='*')],
              [Sg.Button('Ok'), Sg.Text('', key='Err_text', text_color='violetred4', font='Helvetica 15 bold')]]
    window = Sg.Window("KeyMaster - вход", layout, return_keyboard_events=True)
    while True:
        event, values = window.read()
        if event == Sg.WIN_CLOSED:
            break
        if event == 'Ok' or event in ('\r', qt_enter_key1, qt_enter_key2, 'Return:36', 'KP_Enter:104'):
            input_word = values[0]
            if len(input_word) >= 4:
                password = sha256(input_word.encode()).hexdigest()
                tmp = LoginCard(encrypt_key=password, row_id=1)
                if tmp.is_valid():
                    status = True
                    break
                else:
                    attempts -= 1
                    if attempts > 0:
                        err = 'Неверный пароль,\nосталось попыток: ' + str(attempts)
                        window.Element('Err_text').Update(err)
                    else:
                        break
            else:
                attempts -= 1
                if attempts > 0:
                    err = 'Пароль минимум 4 символа.\nосталось попыток: ' + str(attempts)
                    window.Element('Err_text').Update(err)
                else:
                    break

    window.close()
    return status, password


# окно сообщения об ошибках
def err_msg(message):
    # noinspection SpellCheckingInspection
    layout = [[Sg.Text('Oшибка:')],
              [Sg.Text(message)],
              [Sg.Button('Ok')]]
    window = Sg.Window("KeyMaster - ошибка", layout, modal=True)
    while True:
        event, values = window.read()
        if event == 'Ok' or event == Sg.WIN_CLOSED:
            break

    window.close()


# окно подтверждения удаления
def confirm_delete(card_name):
    confirm_flag = False
    layout = [[Sg.Text('Вы уверены, что хотите удалить карточку:')],
              [Sg.Text(card_name + '?')],
              [Sg.Text('Это необратимое действие!!!', text_color='violetred4', font='Helvetica 15 bold')],
              [Sg.Button('Ok'), Sg.Button('Отменить', focus=True)]]
    window = Sg.Window("KeyMaster - подтверждение", layout, modal=True)
    while True:
        event, values = window.read()
        if event == 'Отменить' or event == Sg.WIN_CLOSED:
            break
        if event == 'Ok':
            confirm_flag = True
            break

    window.close()
    return confirm_flag


# окно создания или редактирования карточки логина
def add_or_edit_card(used_list, popular_list, mode, name='', login='', password=''):
    card_added = False
    card_name = name
    card_login = login
    card_password = password
    layout = [[Sg.Text('Имя:')],
              [Sg.InputText(card_name, size=(42, 1), border_width=3)],
              [Sg.Text('Логин:')],
              [Sg.InputCombo(popular_list, default_value=card_login, size=(41, 1))],
              [Sg.Text('Пароль:')],
              [Sg.InputText(card_password, size=(42, 1), border_width=3, key='pass')],
              [Sg.Checkbox('Использовать спецсимволы в пароле', key='special')],
              [Sg.Button('Ok'), Sg.Button('Отменить'), Sg.Button('Сгенерировать пароль')]]
    if mode == 0:
        window = Sg.Window("KeyMaster - новая карточка", layout, modal=True)
    else:
        window = Sg.Window("KeyMaster - редактор карточки", layout, modal=True)
    while True:
        event, values = window.read()
        if event == 'Отменить' or event == Sg.WIN_CLOSED:
            break
        if event == 'Ok':
            error_flag = False
            card_name = values[0]
            card_login = values[1]
            card_password = values['pass']
            if card_name in used_list and mode == 0:
                err_msg('такое имя уже используется.')
                error_flag = True
            else:
                if card_name == '':
                    err_msg('введите имя.')
                    error_flag = True
                if card_login == '' and not error_flag:
                    err_msg('введите логин.')
                    error_flag = True
                if card_password == '' and not error_flag:
                    err_msg('введите пароль.')
                    error_flag = True
            if not error_flag:
                card_added = True
                break
        if event == 'Сгенерировать пароль':
            new_pass = pass_generator(values['special'])
            window.Element('pass').Update(new_pass)

    window.close()
    return card_added, card_name, card_login, card_password


# главное окно
def main():
    global cc1, cc2, cc3

    keys_disabled = True
    my_cards = AllCards(master_key)
    valid, error_numbers = my_cards.validate()
    if not valid:
        message = 'Внимание, карточки ## '
        for number in error_numbers:
            message += str(number) + ', '
        message = message[0:len(message) - 2]
        message += ' повреждены! Используйте бэкап.'
        err_msg(message)
        return

    layout = [[Sg.Text('Главное окно')],
              [Sg.Listbox(my_cards.get_names_list(), select_mode='LISTBOX_SELECT_MODE_SINGLE', enable_events=True,
                          size=(20, 10), key='list'), Sg.Text(key='card')],
              [Sg.Button('Добавить'), Sg.Button('Редактировать', disabled=True, key='edit'),
               Sg.Button('Скопировать пароль', disabled=True, key='copy')],
              [Sg.Button('Удалить', disabled=True, key='delete'), Sg.Text('                               '),
               Sg.Button('Скопировать логин', disabled=True, key='copy_login')]]
    window = Sg.Window("KeyMaster - главная", layout)

    logging.info(
        f'\nзагружено {cc3} карточек.\nвремя чтения {cc1} мкс.\nвремя расшифровки {round(cc2 / 1000000, 2)} секунд.\n'
        'программа запущена.')
    while True:
        event, values = window.read()
        if event == Sg.WIN_CLOSED:
            break

        name = window.Element('list').get()
        if len(name) > 0:
            if not keys_disabled:
                window.Element('edit').update(disabled=True)
                window.Element('copy').update(disabled=True)
                window.Element('copy_login').update(disabled=True)
                window.Element('delete').update(disabled=True)
                keys_disabled = True

        if event == 'Добавить':
            added, name, login, password = add_or_edit_card(my_cards.get_names_list(), my_cards.get_popular_logins(),
                                                            mode=0)
            if added:
                my_cards.add_card(name=name, login=login, password=password)
                my_cards.sort_cards()
                window.Element('list').Update(my_cards.get_names_list())

        if event == 'list':
            name = window.Element('list').get()
            if len(name) > 0:
                if keys_disabled:
                    window.Element('edit').update(disabled=False)
                    window.Element('copy').update(disabled=False)
                    window.Element('copy_login').update(disabled=False)
                    window.Element('delete').update(disabled=False)
                    keys_disabled = False

                card = my_cards.get_card_by_name(name[0])
                window.Element('card').Update(card)

        if event == 'copy':
            name = window.Element('list').get()
            if len(name) > 0:
                password = my_cards.get_card_password(name[0])
                Sg.clipboard_set(password)

        if event == 'copy_login':
            name = window.Element('list').get()
            if len(name) > 0:
                login = my_cards.get_card_login(name[0])
                Sg.clipboard_set(login)

        if event == 'edit':
            name = window.Element('list').get()
            if len(name) > 0:
                card_name = name[0]
                card_id = my_cards.get_card_id(card_name)
                login = my_cards.get_card_login(card_name)
                password = my_cards.get_card_password(card_name)
                added, name_new, login_new, password_new = add_or_edit_card(my_cards.get_names_list(),
                                                                            my_cards.get_popular_logins(), mode=1,
                                                                            name=card_name, login=login,
                                                                            password=password)
                if added:
                    my_cards.update_card(row_id=card_id, name=name_new, login=login_new, password=password_new)
                    my_cards.sort_cards()
                    window.Element('list').Update(my_cards.get_names_list())

        if event == 'delete':
            name = window.Element('list').get()
            if len(name) > 0:
                card_name = name[0]
                confirmed = confirm_delete(card_name)
                if confirmed:
                    my_cards.delete_card(card_name)
                    my_cards.sort_cards()
                    window.Element('list').Update(my_cards.get_names_list())

    window.close()
    logging.info('программа завершена.')


# начало исполнения кода.
if __name__ == '__main__':
    kill(getppid(), SIGKILL)

    if not path.exists('base.sqlite'):
        login_status, master_key = system_init()
        if login_status:
            con = sqlite3.connect('base.sqlite')
            cursor = con.cursor()
            cursor.executescript('''
                        create table logins
                            (
                                row_id integer
                                    constraint logins_pk
                                        primary key autoincrement,
                                row_name text not null,
                                login text not null,
                                password text not null,
                                hash text not null                            
                            );
                        ''')
            con.commit()
            pass_save = LoginCard(master_key, name='master_login', login='0', password='0')
    else:
        con = sqlite3.connect('base.sqlite')
        cursor = con.cursor()
        logging.info('соединение с бд установлено, ожидаем пароль...')
        login_status, master_key = system_login()

    if login_status:
        logging.info('логин - ок. запуск...')
        main()