#!/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)


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')