# 20201215 王馨瑶
# -*- coding: UTF-8 -*-
# 文件名：clien1.py

import numpy
import os
import time
import psutil
import base64
from Crypto.Cipher import AES
import socket

BLOCK_SIZE = 16  # Bytes
pad = lambda s: s + (BLOCK_SIZE - len(s) % BLOCK_SIZE) * \
                chr(BLOCK_SIZE - len(s) % BLOCK_SIZE)

def txt(name, text):  # 定义函数名
    b = os.getcwd()[:-4] + 'new\\'
    if not os.path.exists(b):  # 判断当前路径是否存在，没有则创建new文件夹
        os.makedirs(b)
      # 在当前py文件所在路径下的new文件中创建txt
    with open(name,'w')as file1:
        file1.write(text)  # 写入内容信息
        file1.close()
    print('ok')

def add_to_16(value):  # str不是16的倍数那就补足为16的倍数
    while len(value) % 16 != 0:
        value += '\0'
    return str.encode(value)  # 返回bytes
'''
def encrypt(filename,af_filename):#加密方法
    key = input("Please set a key:")
    text = open(filename,'rb').read()
    open(filename,'rb').close()
    text = str(text)
    aes = AES.new(add_to_16(key), AES.MODE_ECB)
    encrypt_aes = aes.encrypt(add_to_16(text))
    encrypted_text = str(base64.encodebytes(encrypt_aes), encoding='utf-8')
    txt(af_filename, encrypted_text)
    time.sleep(1)  # 推迟执行
    print('... File encrypted successfully ...')
'''

def encrypt(filename,af_filename):
    key = input("Please set a key:")
    data = str(open(filename, 'rb').read())
    open(filename, 'rb').close()
    key = key.encode('utf8')
    data = pad(data)
    cipher = AES.new(key, AES.MODE_ECB)
    # 加密后得到的是bytes类型的数据，使用Base64进行编码,返回byte字符串
    result = cipher.encrypt(data.encode())
    encodestrs = base64.b64encode(result)
    enctext = encodestrs.decode('utf8')
    txt(af_filename, enctext)
    time.sleep(1)
    print('... File encrypted successfully ...')
    return enctext



def send_out(af_filename):
    with open(af_filename, 'rb') as file:
        i = file.read()
        amanda.send(i)
    # amanda.send('end'.encode())
    time.sleep(1)  # 推迟执行
    print("... File sent successfully ...")

#
# ip = input("pleas input the serve ip:")
# port = int(input("please input the port:"))
amanda = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
amanda.connect(('127.0.0.1', 8001))  # 建立连接，参数必须是元组的形式
print(amanda.recv(1024))  # 接收welcome！
# str= input("please input your message:")
# amanda.sendall(str.encode())
# data = s.recv(1024)
# print (data.decode())
print("Hello!I'm Amanda.\nGlad to see you!\nIf you want to encrypt a file,enter \"encrypt file\".\nTo end the chat,"
      "enter \"end\"")

# def f_send(f_name):

while True:
    fun = input("Input your choice:")
    amanda.sendall(fun.encode())
    if fun == 'end':
        print("Ok,Bye!")
        break
    elif fun == 'encrypt file':
        filename = input("Please Input the file name which you want to encrypt:")
        with open(filename, 'w')as file:
            content = input("please input the content of " + filename + ':')
            try:
                file.writelines(content)
            finally:
                file.close()
        af_filename = 'encrypted.txt'
        amanda.sendall(af_filename.encode())
        # amanda.sendall('amanda'.encode())
        encrypt(filename, af_filename)
        send_out(af_filename)

    amanda.close()