#!/usr/bin/env python3
import sys
import hashlib
import random
import re
import os
import shutil
import argparse

flagtext = """
   Insert required amount of spaces *here* to decrypt the message.

   The flag is:

       """
plaintext = ""

# parse command line
parser = argparse.ArgumentParser()
parser.add_argument('-n', dest="tasks", required=False, default=1, type=int, help="Number of tasks to create")
parser.add_argument('-m', metavar="FILE", dest="txtfile", required=True, help="File with plaintext")
parser.add_argument('fname', metavar="DIR", nargs=1, help="Output directory")
args = parser.parse_args()

# create output directory
os.mkdir(args.fname[0])

# read the plaintext file
with open(args.txtfile, "r") as f:
    plaintext = f.read()
    if len(plaintext) < len(flagtext)+64:
        print("Plaintext too short,",len(flagtext)+64,"bytes required")
        sys.exit(1)

# repeat for each participant
for i in range(1,args.tasks+1):
    
    dest = args.fname[0] + "/"

    # make subdirectory if more than 1 task
    if args.tasks > 1:
        dest += str(i)
        os.mkdir(dest)

    # generate random flag
    h = hashlib.sha1()
    h.update(bytes(str(random.uniform(714873,98734018430713)), "ascii"))
    flag = "ECSC{" + h.hexdigest() + "}"
    print("flag for participant",i,":",flag)

    # generate pad table and location for flag
    padtable = [int(random.uniform(0,255)) for x in range(len(plaintext))]
    m = re.findall(r'(.*\*)(here)(\*.*)', flagtext, re.DOTALL)
    offset = int(random.uniform(len(m[0][1]), len(plaintext)-len(flagtext)-len(flag)))

    # fix pad table to get correct initial message
    for x in range(len(m[0][0]),len(flagtext)):
        padtable[x] = padtable[x+offset]

    # create encrypted messages
    loc = len(flagtext) + offset
    encflag = [(ord(a) + b)%256 for a, b in zip(flag, padtable[loc:loc+len(flag)])]
    with open(dest+"/mainfile.enc", "wb") as f:
        f.write(bytes([(ord(a) + b)%256 for a, b in zip(flagtext, padtable[0:len(flagtext)])]))
        f.write(bytes(encflag))

    with open(dest+"/"+args.txtfile+".enc", "wb") as f:
        f.write(bytes([(ord(a) + b)%256 for a, b in zip(plaintext, padtable)]))

    # copy the plaintext to destination
    shutil.copy2(args.txtfile, dest)