# Odbiór wiadomości
import hashlib, random, os, time
from binascii import hexlify
from socket import *
import r09_telnet as ct

serverDH = ct.gen_server_DH()
serverSecret = 0

def get_dh_sharedsecret():
    key = int(open('klucz_publiczny_dh_klient.pem').read())
    serverDH.generateSharedKey(key)
    serverDH.getSharedKey()

    serverDH.generateSharedKey(key)
    return  (serverDH.sharedSecret)

def get_dh_sharedkey():
    key = int(open('klucz_publiczny_dh_klient.pem').read())
    serverDH.generateSharedKey(key)
    serverDH.getSharedKey()

    serverDH.generateSharedKey(key)

    private_key = serverDH.key

    return  private_key

def decrypt(ciphertext, usePKI, useDH, serverSecret):
    try:
        msg = ct.decrypt(ciphertext, usePKI, useDH, serverSecret)
    except:
        msg = ciphertext
    return msg

def main(): 
    # Zmienne pozwalające wybrać schemat szyfrowania
    useClientPKI = False
    useDHKey = False
    serverSecret = 0

    # Zmienne związane z pracą serwera
    key = ""
    host = ""
    port = 8080
    buf = 1024 * 2
    addr = (host, port)
    UDPSock = socket(AF_INET, SOCK_DGRAM)
    UDPSock.bind(addr)

    # Wiadomość powitalna
    print("Oczekuję na wiadomość...")

    # Pętla nasłuchująca
    while True:
        # Wczytanie danych otrzymanych od klienta
        (data, addr) = UDPSock.recvfrom(buf)
        
        # Przesłanie danych do odszyfrowania
        plaintext = decrypt(data, useClientPKI, useDHKey, serverSecret)

        # Sprawdzenie, czy klient nie użył polecenia, np. addPKI lub addDH, jeżeli użył to wiadomość jest zastępowana komunikatem
        result = ct.check_server_command(plaintext)

        if result == 10: # Brak wiadomości, wyłączenie szyfrowania
            plaintext = 'Szyfrowanie PKI wyłączone!'
            msg = plaintext
        elif result == 11: # Włączenie szyfrowania
            plaintext = 'Szyfrowanie PKI włączone!'
            msg = plaintext
        elif result == 20: # Wyłączenie dh
            clientKey = plaintext
            plaintext = 'Protokół Diffiego-Hellmana nie aktywny!'
            msg = plaintext
        elif result == 21: # Właczenie dh 
            plaintext = 'Protokół Diffiego-Hellmana aktywny!'
            msg = plaintext
        else:
            # Jeżeli nie zastąpiono wiadomości komunikatem, to należy ją odkodować
            msg = str(plaintext, 'utf-8')
        
        # Przetwarzanie poleceń
        if result == 0:
            # Brak szyfrowania
            break
        if result == 10:
            # Wyłączenie PKI
            useClientPKI = False
        if result == 11:
            # Włączenie PKI
            useClientPKI = True
            # Informacja o znalezieniu certyfikatów
            print("Certyfikat klienta został znaleziony ...")
        if result == 20:
            # Wyłączenie dh
            useDHKey = False
        if result == 21:
            # Włączenie dh
            useDHKey = True
            print("Protokół wymiany kluczy Diffiego-Hellmana jest aktywyny ...")            
            serverSecret = get_dh_sharedkey()

        # Jeżeli szyfrowanie jest aktywne, otrzymana wiadomość była zaszyfrowana
        if useClientPKI == True or useDHKey == True:
            print("Otrzymałem zaszyfrowaną wiadomość: " + msg)
        else:
            print("Otrzymałem wiadomość: " + msg)
        


    UDPSock.close()
    os._exit(0) 

if __name__ == '__main__': 
    main() 


