lab5

import sys
from PyQt5.Qt import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
import random
import base64
import pickle
from Crypto import Random
from Crypto.Cipher import AES
import rsa
import hashlib
from ocb.aes import AES as AES1
from ocb import OCB

class NotRealNetwork:

    '''буфер, симулирующий сеть'''

    def __init__(self):
        self.contact=dict()#здесь лежит информация про зарегистрированные объекты и их "почтовые ящики"
        # потоки будут по очереди использовать этот словарь через блокировку

    def enter(self, name):

        '''Вход объекта в "сеть". Через имя объект может получить информацию о том,
        что с ним хотят свяаться. Тогда в списке будет кортеж их имени отправителя и
        содержания послания'''

        if self.contact.get(name) is None:
            self.contact[name]=[]
            return True
        else:
            return False

    def delete(self, name):
        if self.contact.get(name) is None:
            return False
        else:
            self.contact.pop(name)
            return True


    def write(self, send_name, rec_name, data):

        '''Пишет в ящик получателя информацию и имя отправителя'''
        if self.contact.get(rec_name) is None:
            return False
        else:
            self.contact[rec_name].append((send_name, data))
            return True

    def read(self, name):

        '''Выдаёт получателю следующий по очереди отправленный ему пакет, если он есть'''

        if self.contact.get(name):
            return self.contact[name].pop(0)
        else:
            return []

class AESCipher:

    def __init__( self, key, mode):
        self.key = self.set_key(key)
        self.mode=mode

    def set_key(self, key):
        key=str(key).encode('utf-8')
        st=16-len(key)%16
        st1=b'1'*st
        key=key+st1
        return key

    def encrypt(self, raw):
        BS = 16
        pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS)##
        pad=lambda s:s+(BS-len(s)%BS)*chr(BS-len(s)%BS).encode('utf-8')#b'\xd0'
        unpad = lambda s : s[0:-s[-1]]
        raw = pad(raw)
        iv = Random.new().read(AES.block_size)
        if self.mode=='ecb':
            cipher = AES.new(self.key, AES.MODE_ECB, iv)
        elif self.mode=='cbc':
            cipher = AES.new(self.key, AES.MODE_CBC, iv)
        elif self.mode=='ofb':
            cipher = AES.new(self.key, AES.MODE_OFB, iv)
        elif self.mode=='cfb':
            cipher = AES.new(self.key, AES.MODE_CFB, iv)
        #elif self.mode=='ocb':
        #    cipher = AES.new(self.key, AES.MODE_OCB, iv)
        return base64.b64encode(iv + cipher.encrypt(raw))

    def decrypt(self, enc):
        BS = 16
        #pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS)
        unpad = lambda s : s[0:-s[-1]]
        enc = base64.b64decode(enc)
        iv = enc[:16]
        #cipher = AES.new(self.key, AES.MODE_ECB, iv)
        if self.mode=='ecb':
            cipher = AES.new(self.key, AES.MODE_ECB, iv)
        elif self.mode=='cbc':
            cipher = AES.new(self.key, AES.MODE_CBC, iv)
        elif self.mode=='ofb':
            cipher = AES.new(self.key, AES.MODE_OFB, iv)
        elif self.mode=='cfb':
            cipher = AES.new(self.key, AES.MODE_CFB, iv)
        return unpad(cipher.decrypt( enc[16:] ))

class Subject:

    '''Субьект протокола, управляется пользователем'''

    def __init__(self, name, web):
        self.name=name
        self.web=web
        self.asim_reg_database=dict()
        self.sim_reg_database=dict()

    def reg_web(self):
        f=self.web.enter(self.name)
        return f

    def check(self, system_name, suc, fail):

        '''проверяет результаты действия. suc - строка, возврат которо обозначает успех, fail - наоборот'''
        obj=objs=self.web.read(self.name) #берём первое письмо и запоминаем
        if not obj or obj is None: #если ничего не написали, то следующая попытка
            return False

        while not (obj[0]==system_name and obj[1]['nomination']==suc): #другое письмо
            if obj[0]==system_name and obj[1]['nomination']==fail: #пришёл отказ - регистрация не удалать
                return False

            self.web.write(obj[0], self.name, obj[1]) #возвращаем ненужное письмо обратно в очередь
            obj=self.web.read(self.name) #и берём следующее
            if obj[1] is objs[1]: #если все письма уже перебрали, а нужного не нашли
                self.web.write(obj[0], self.name, obj[1])  # возвращаем ненужное письмо обратно в очередь
                break

        else: # нашли нужное письмо!
            return obj

    def asim_reg(self, name, other_name):
        (pubkey, privkey) = rsa.newkeys(512)
        self.privkey=privkey
        self.pubkey=pubkey
        send = dict()
        send['nomination'] = 'asim reg'
        send['name'] = name
        send['password'] = self.pubkey
        f = self.web.write(self.name, other_name, send)
        return f

    def asim_reg_check(self, name, other_name):
        send = dict()
        res = self.check(other_name, 'asim reg', 'asim fail')
        if not res or self.asim_reg_database.get(res[1]['name']):
            send['nomination']='asim reg fail'
            self.web.write(name, res[1]['name'], send)
            return False

        self.asim_reg_database[res[1]['name']]=res[1]['password']

        send['nomination'] = 'asim reg done'
        f = self.web.write(name, res[1]['name'], send)
        return (f, res[1]['name'], res[1]['password'])

    def asim_reg_result(self, other_name):
        res=self.check(other_name, 'asim reg done', 'asim reg fail')
        return res

    def hyb_inquiry(self, message, name, other_name, mode):
        send = dict()
        send['nomination'] = 'hyb inquiry'
        send['name'] = name
        send['other name'] = other_name

        ses_key = random.randint(1000000000000000, 9999999999999999)
        self.ses_key=ses_key
        crypto=pickle.dumps(ses_key, 2)

        if self.asim_reg_database.get(other_name) is None:
            return False

        crypto=rsa.encrypt(crypto, self.asim_reg_database[other_name])
        send['ses_key']=crypto

        self.message=message
        cipher = AESCipher(self.ses_key, mode)
        message=pickle.dumps((self.message), 2)
        message = cipher.encrypt(message)
        send['message'] = message

        f = self.web.write(self.name, other_name, send)
        return (f, self.message, ses_key, message)

    def hyb_res(self, other_name, mode):
        res=self.check(other_name, 'hyb inquiry', 'hyb fail')
        if not res:
            return False

        ses_key = rsa.decrypt(res[1]['ses_key'], self.privkey)
        ses_key = pickle.loads(ses_key)
        self.ses_key = ses_key

        cipher = AESCipher(self.ses_key, mode)
        decrypted=cipher.decrypt(res[1]['message'])
        decrypted = pickle.loads(decrypted)
        return (decrypted, ses_key, res[1]['message'])

    def asim_inquiry(self, message, name, other_name, mode):
        send = dict()
        send['nomination'] = 'asim inquiry'
        send['name'] = name
        send['other name'] = other_name

        ses_key = random.randint(1000000000000000, 9999999999999999)
        self.ses_key = ses_key
        crypto = pickle.dumps(ses_key, 2)

        if self.asim_reg_database.get(other_name) is None:
            return False

        crypto = rsa.encrypt(crypto, self.asim_reg_database[other_name])
        send['ses_key'] = crypto

        self.message = message
        cipher = AESCipher(self.ses_key, mode)
        message = pickle.dumps((self.message), 2)

        sign=hashlib.sha1(message)
        sign.update(pickle.dumps((self.asim_reg_database[other_name]),2))
        sign=sign.hexdigest()
        send['sign']=sign

        message = cipher.encrypt(message)
        send['message'] = message



        f = self.web.write(self.name, other_name, send)
        return (f, self.message, ses_key, message, sign)

    def asim_res(self, other_name, mode):
        res=self.check(other_name, 'asim inquiry', 'asim fail')
        if not res:
            return False

        ses_key = rsa.decrypt(res[1]['ses_key'], self.privkey)
        ses_key = pickle.loads(ses_key)
        self.ses_key = ses_key

        cipher = AESCipher(self.ses_key, mode)
        decrypted = cipher.decrypt(res[1]['message'])
        sign = hashlib.sha1(decrypted)
        sign.update(pickle.dumps((self.pubkey), 2))
        sign=sign.hexdigest()
        if sign != res[1]['sign']:
            return False

        decrypted = pickle.loads(decrypted)

        return (decrypted, ses_key, res[1]['message'], res[1]['sign'])

    def asim_ocb_inquiry(self, message, name, other_name):
        send = dict()
        send['nomination'] = 'asim ocb inquiry'
        send['name'] = name
        send['other name'] = other_name

        ses_key = bytearray().fromhex('A45F5FDEA5C088D1D7C8BE37CABC8C5C')
        self.ses_key = ses_key
        nonce = bytearray(range(16))

        if self.asim_reg_database.get(other_name) is None:
            return False

        send['ses_key'] = (ses_key, nonce)

        self.message = message
        aes = AES1(128)
        ocb = OCB(aes)
        ocb.setKey(ses_key)
        ocb.setNonce(nonce)
        plaintext = bytearray(bytes(message, encoding = 'utf-8'))
        (tag, ciphertext) = ocb.encrypt(plaintext, bytearray(b'123'))

        send['message'] = ciphertext
        send['tag']=tag

        f = self.web.write(self.name, other_name, send)
        return (f, self.message, ses_key, ciphertext)

    def asim_ocb_res(self, other_name):
        res=self.check(other_name, 'asim ocb inquiry', 'asim ocb fail')
        if not res:
            return False

        (ses_key, nonce) = res[1]['ses_key']
        aes = AES1(128)
        ocb = OCB(aes)
        ocb.setKey(ses_key)
        ocb.setNonce(nonce)
        (is_authentic, plaintext2) = ocb.decrypt(bytearray(b'123'), res[1]['message'], res[1]['tag'],)
        decrypted=plaintext2.decode("utf-8")

        if not is_authentic:
            return False

        return (decrypted, ses_key, res[1]['message'])

class SubjectDraw(QWidget):

    '''Визуализация субъектов'''

    def __init__(self, web, nm):
        QWidget.__init__(self)
        self.initUI(nm)
        self.subject=dict() # массив с зарегистрированными субъектами
        self.web=web

    def initUI(self, nm):
        self.resize(500, 700)
        vbox = QVBoxLayout()

        hbox1 = QHBoxLayout()
        self.name = QLabel("Имя")
        hbox1.addWidget(self.name)
        self.message_name = QLabel("Сообщение")
        hbox1.addWidget(self.message_name)

        self.other_name = QLabel("Другое имя")
        hbox1.addWidget(self.other_name)
        self.password_label = QLabel("Пароль")
        hbox1.addWidget(self.password_label)
        vbox.addLayout(hbox1)

        hbox2 = QHBoxLayout()
        self.name_edit = QLineEdit()
        self.name_edit.setText("name"+nm)
        hbox2.addWidget(self.name_edit)
        self.message_edit = QLineEdit()
        self.message_edit.setText("mess")
        hbox2.addWidget(self.message_edit)
        self.other_name_edit = QLineEdit()

        if nm=='A':
            self.other_name_edit.setText('nameB')
        elif nm=='B':
            self.other_name_edit.setText('nameA')

        hbox2.addWidget(self.other_name_edit)
        self.password_edit = QLineEdit()
        self.password_edit.setText(str(ord(nm)*10000))
        hbox2.addWidget(self.password_edit)
        vbox.addLayout(hbox2)

        self.plain_text = QPlainTextEdit()
        vbox.addWidget(self.plain_text)

        hbox = QHBoxLayout()
        self.create_button = QPushButton("Создать объект")
        self.create_button.clicked.connect(self.create_subject)
        hbox.addWidget(self.create_button)
        self.web_button = QPushButton("Зайти в сеть")
        self.web_button.clicked.connect(self.web_reg)
        hbox.addWidget(self.web_button)
        self.asim_reg_button = QPushButton("Ассиметричная регистрация")
        self.asim_reg_button.clicked.connect(self.asim_reg)
        hbox.addWidget(self.asim_reg_button)
        self.asim_check_button = QPushButton("Проверить чужую регистрацию")
        self.asim_check_button.clicked.connect(self.asim_check)
        hbox.addWidget(self.asim_check_button)
        self.asim_res_button = QPushButton("Результат своей регистрации")
        self.asim_res_button.clicked.connect(self.asim_reg_result)
        hbox.addWidget(self.asim_res_button)
        vbox.addLayout(hbox)

        hbox = QHBoxLayout()
        self.hyb_inquiry_button = QPushButton("Послать сообщение ECB")
        self.hyb_inquiry_button.clicked.connect(self.ecb_hyb_inquiry)
        hbox.addWidget(self.hyb_inquiry_button)
        self.hyb_res_button = QPushButton("Получить сообщение ECB")
        self.hyb_res_button.clicked.connect(self.ecb_hyb_res)
        hbox.addWidget(self.hyb_res_button)
        self.hyb_inquiry_button1 = QPushButton("Послать сообщение CBC")
        self.hyb_inquiry_button1.clicked.connect(self.cbc_hyb_inquiry)
        hbox.addWidget(self.hyb_inquiry_button1)
        self.hyb_res_button1 = QPushButton("Получить сообщение CBC")
        self.hyb_res_button1.clicked.connect(self.cbc_hyb_res)
        hbox.addWidget(self.hyb_res_button1)
        self.hyb_inquiry_button2 = QPushButton("Послать сообщение OFB")
        self.hyb_inquiry_button2.clicked.connect(self.ofb_hyb_inquiry)
        hbox.addWidget(self.hyb_inquiry_button2)
        self.hyb_res_button2 = QPushButton("Получить сообщение OFB")
        self.hyb_res_button2.clicked.connect(self.ofb_hyb_res)
        hbox.addWidget(self.hyb_res_button2)
        self.hyb_inquiry_button3 = QPushButton("Послать сообщение CFB")
        self.hyb_inquiry_button3.clicked.connect(self.cfb_hyb_inquiry)
        hbox.addWidget(self.hyb_inquiry_button3)
        self.hyb_res_button3 = QPushButton("Получить сообщение CFB")
        self.hyb_res_button3.clicked.connect(self.cfb_hyb_res)
        hbox.addWidget(self.hyb_res_button3)
        vbox.addLayout(hbox)

        hbox1 = QHBoxLayout()
        self.asim_inquiry_button = QPushButton("Послать сообщение EM")
        self.asim_inquiry_button.clicked.connect(self.asim_inquiry)
        hbox1.addWidget(self.asim_inquiry_button)
        self.asim_res_button = QPushButton("Получить сообщение EM")
        self.asim_res_button.clicked.connect(self.asim_res)
        hbox1.addWidget(self.asim_res_button)
        self.asim_ocb_inquiry_button = QPushButton("Послать сообщение OCB")
        self.asim_ocb_inquiry_button.clicked.connect(self.asim_ocb_inquiry)
        hbox1.addWidget(self.asim_ocb_inquiry_button)
        self.asim_ocb_res_button = QPushButton("Получить сообщение OCB")
        self.asim_ocb_res_button.clicked.connect(self.asim_ocb_res)
        hbox1.addWidget(self.asim_ocb_res_button)
        vbox.addLayout(hbox1)

        self.setLayout(vbox)

    def create_subject(self):
        name=self.name_edit.text()
        if not self.subject.get(name) is None:
            self.plain_text.appendPlainText("Уже зарегистрирован")
        else:
            self.subject[name]=Subject(name, self.web)
            self.plain_text.appendPlainText("Создан")

    def web_reg(self):
        name = self.name_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Объекта нет")
        else:
            if self.subject[name].reg_web():
                self.plain_text.appendPlainText("Зашёл")
            else:
                self.plain_text.appendPlainText("Вход этим именем уже произведён")

    def asim_reg(self):
        name = self.name_edit.text()
        password=self.password_edit.text()
        other_name = self.other_name_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f=self.subject[name].asim_reg(name, other_name)
            if f:
                self.plain_text.appendPlainText("Запрос отправлен")
            else:
                self.plain_text.appendPlainText("Ошибка")

    def asim_check(self):
        name = self.name_edit.text()
        other_name = self.other_name_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f = self.subject[name].asim_reg_check(name, other_name)
            if f:
                self.plain_text.appendPlainText("Прислан запрос от " + str(f[1]) + "с паролем " + str(f[2]))
            else:
                self.plain_text.appendPlainText("Ошибка")

    def asim_reg_result(self):
        name = self.name_edit.text()
        other_name = self.other_name_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f = self.subject[name].asim_reg_result(other_name)
            if f:
                self.plain_text.appendPlainText("Регистрация успешна")
            else:
                self.plain_text.appendPlainText("Ошибка")

    def hyb_inquiry(self):
        name = self.name_edit.text()
        other_name = self.other_name_edit.text()
        message = self.message_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f=self.subject[name].hyb_inquiry(message, name, other_name, self.mode)
            if f:
                self.plain_text.appendPlainText("Запрос отправлен. Послано "+str(f[1])+', зашифрованное на ключе '+str(f[2])+". Результат: "+str(f[3]))
            else:
                self.plain_text.appendPlainText("Ошибка")

    def hyb_res(self):
        name = self.name_edit.text()
        other_name = self.other_name_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f=self.subject[name].hyb_res(other_name, self.mode)
            if f:
                self.plain_text.appendPlainText("Пришло сообщение "+str(f[0])+', расшифрованное на ключе '+str(f[1])+" из шифртекста "+str(f[2]))
            else:
                self.plain_text.appendPlainText("Ошибка")

    def ecb_hyb_inquiry(self):
        self.mode='ecb'
        self.hyb_inquiry()

    def ecb_hyb_res(self):
        self.mode='ecb'
        self.hyb_res()

    def cbc_hyb_inquiry(self):
        self.mode = 'cbc'
        self.hyb_inquiry()

    def cbc_hyb_res(self):
        self.mode = 'cbc'
        self.hyb_res()

    def ofb_hyb_inquiry(self):
        self.mode = 'ofb'
        self.hyb_inquiry()

    def ofb_hyb_res(self):
        self.mode = 'ofb'
        self.hyb_res()

    def cfb_hyb_inquiry(self):
        self.mode = 'cfb'
        self.hyb_inquiry()

    def cfb_hyb_res(self):
        self.mode = 'cfb'
        self.hyb_res()

    def asim_inquiry(self):
        name = self.name_edit.text()
        other_name = self.other_name_edit.text()
        message = self.message_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f=self.subject[name].asim_inquiry(message, name, other_name, 'ecb')
            if f:
                self.plain_text.appendPlainText("Запрос отправлен. Послано  "+str(f[1])+', зашифрованное на ключе '+str(f[2])+" в виде "+str(f[3])+". Подпись: "+str(f[4]))
            else:
                self.plain_text.appendPlainText("Ошибка")

    def asim_res(self):
        name = self.name_edit.text()
        other_name = self.other_name_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f=self.subject[name].asim_res(other_name, 'ecb')
            if f:
                self.plain_text.appendPlainText("Пришло "+str(f[0])+', расшифрованное ключом '+str(f[1])+" из шифртекста "+str(f[2])+" с подписью "+str(f[3]))
            else:
                self.plain_text.appendPlainText("Ошибка")

    def asim_ocb_inquiry(self):
        name = self.name_edit.text()
        other_name = self.other_name_edit.text()
        message = self.message_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f=self.subject[name].asim_ocb_inquiry(message, name, other_name)
            if f:
                self.plain_text.appendPlainText("Запрос отправлен. Послано  "+str(f[1])+', зашифрованное на ключе '+str(f[2])+" в виде "+str(f[3]))
            else:
                self.plain_text.appendPlainText("Ошибка")

    def asim_ocb_res(self):
        name = self.name_edit.text()
        other_name = self.other_name_edit.text()
        if self.subject.get(name) is None:
            self.plain_text.appendPlainText("Не зарегистрирован")
            return
        else:
            f=self.subject[name].asim_ocb_res(other_name)
            if f:
                self.plain_text.appendPlainText("Пришло "+str(f[0])+', расшифрованное ключом '+str(f[1])+" из шифртекста "+str(f[2]))
            else:
                self.plain_text.appendPlainText("Ошибка")

w=NotRealNetwork()
app = QApplication(sys.argv)
s1 = SubjectDraw(w, 'A')
s2 = SubjectDraw(w, 'B')
s1.show()
s2.show()

app.exit(app.exec_())