import socket, sys
import rsa
from aes import *


# 读取文件信息
def read_message(filename):
    msg = ""
    with open(filename) as f:
        try:
            for line in f.readlines():
                msg += line
        finally:
            f.close()
    return msg


# 写入文件信息
def write_message(filename, msg):
    try:
        f = open(filename, 'a')
        f.write(msg + '\n')
    finally:
        f.close()


def receiver():
    # 设置接收端端口
    address = ('127.0.0.1', 6666)

    # 建立接收端
    try:
        receiver = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        print("Receiving terminal start!")
    except socket.error as msg:
        print("Socket failed. Error code: " + str(msg[0]) + ", Error message: " + msg[1])
        sys.exit()

    # 绑定监听
    try:
        receiver.bind(address)
        print("Binding port succeed!")
    except socket.error as msg:
        print("Bind failed. Error code: " + str(msg[0]) + ", Error message: " + msg[1])
        sys.exit()

    # 设置连接数为 1
    receiver.listen(1)

    # 接收连接,读取数据
    while True:
        connect, address = receiver.accept()
        data = connect.recv(4096)
        print('\nReceived messages:\n' + data.decode())
        if not data:
            break
        connect.sendall(b'Receiver: messages have received!')
        return data.decode()

    # 断开连接
    connect.close()
    receiver.close()


# 接收密文，并解密消息
def main():
    # 获取消息
    cipher_message = receiver().split('.')
    aes_key_encrypted = cipher_message[0]
    message_encrypted = cipher_message[1]
    rblk = message_encrypted.split(",")
    rblk = list(map(int, rblk))

    # 获取rsa私钥（接收端用私钥解密）
    rsa_d = int(read_message('./rsa_key/rsa_d.txt'), 16)
    rsa_n = int(read_message('./rsa_key/rsa_n.txt'), 16)

    # 解密得到AES对称密钥
    aes_key = rsa.rsa_decrypt(aes_key_encrypted, rsa_d, rsa_n)
    print("得到AES密钥：" + aes_key)
    key = StringToListN(aes_key)

    # 解密
    b = AESD(rblk, key, 10)
    b.AesDecrpyt()


if __name__ == "__main__":
    main()