#!/usr/bin/env python3

import time
import random
import string
import gpg
import subprocess
from pgpy import PGPMessage, PGPKey, constants
#from pysequoia import Cert, Context
from multiprocessing import Pool
from tqdm import tqdm

N = 5000
K = 16

SIGN = True
SIGN_WITH = ('4A3E34DF1E4569E72A197950BAFB02D14732B480', 'sign.asc')
ENCRYPT_TO = ('5C0EF910C3F5FDB51BCE5630C96DCF4E534983F4', 'encrypt.asc')


def gpg_worker(i):
    with gpg.Context(armor=True) as ctx:
        ctx.set_engine_info(gpg.constants.protocol.OpenPGP,
                            home_dir='gpg' + str(i%8))
        for _ in tqdm(range(N)):
            msg = ''.join([random.choice(string.ascii_lowercase) for _ in range(100)]).encode()
            rkey = list(ctx.keylist(pattern=ENCRYPT_TO[0], secret=False))
            if SIGN:
                skey = list(ctx.keylist(pattern=SIGN_WITH[0], secret=False))
            else:
                skey = None
            ciphertext, result, sign_result = ctx.encrypt(msg, recipients=rkey, sign=skey, always_trust=True, compress=False)

def gpg_singlehome_worker(i):
    gpg_worker(0)

def pgpy_worker(i):
    for _ in tqdm(range(N)):
        msg = PGPMessage.new(''.join([random.choice(string.ascii_lowercase) for _ in range(100)]),
                             compression=constants.CompressionAlgorithm.Uncompressed)
        if SIGN:
            skey, _ = PGPKey.from_file(SIGN_WITH[1])
            msg |= skey.sign(msg)
        rkey, _ = PGPKey.from_file(ENCRYPT_TO[1])
        enc = rkey.encrypt(msg)
        x = str(enc)

def sqop_worker(i):
    for _ in tqdm(range(N)):
        msg = ''.join([random.choice(string.ascii_lowercase) for _ in range(100)])
        if SIGN:
            cmdline = ['/usr/bin/sqop', 'encrypt', '--sign-with', SIGN_WITH[1], ENCRYPT_TO[1]]
        else:
            cmdline = ['/usr/bin/sqop', 'encrypt', ENCRYPT_TO[1]]
        p = subprocess.Popen(cmdline, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        out, _ = p.communicate(msg.encode())

def pysq_worker(i):
    ctx = Context.standard()
    for _ in tqdm(range(N)):
        msg = ''.join([random.choice(string.ascii_lowercase) for _ in range(100)])
        if SIGN:
            skey = Cert.from_file(SIGN_WITH[1])
        else:
            skey = None
        rkey = Cert.from_file(ENCRYPT_TO[1])
        enc = ctx.encrypt(skey, rkey, msg)
        print(enc)


if __name__ == '__main__':
    gpg_start = time.monotonic()
    with Pool(K) as p:
        p.map(gpg_worker, range(K))
    gpg_end = time.monotonic()
    gpg_single_start = time.monotonic()
    with Pool(K) as p:
        p.map(gpg_singlehome_worker, range(K))
    gpg_single_end = time.monotonic()
    pgpy_start = time.monotonic()
    with Pool(K) as p:
        p.map(pgpy_worker, range(K))
    pgpy_end = time.monotonic()
    sqop_start = time.monotonic()
    with Pool(K) as p:
        p.map(sqop_worker, range(K))
    sqop_end = time.monotonic()
    #pysq_start = time.monotonic()
    #with Pool(K) as p:
    #    p.map(pysq_worker, range(K))
    #pysq_end = time.monotonic()
    print(f'GnuPG MH : {gpg_end - gpg_start}s')
    print(f'GnuPG SH : {gpg_single_end - gpg_single_start}s')
    print(f'PGPy     : {pgpy_end - pgpy_start}s')
    print(f'sqop     : {sqop_end - sqop_start}s')
    #print(f'pysequoia: {pysq_end - pysq_start}s')