import re
import subprocess

from test import MailTest


def import_pubkey(key, GNUPG_DIRECTORY):
    with subprocess.Popen(["gpg",
                           "--homedir=" + GNUPG_DIRECTORY,
                           "--batch",
                           "--import",
                           "-"], env={"LANG": 'C'}, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as p:
        p.stdin.write(key)
        p.stdin.close()
        p.wait()
        gpg_output = p.stderr.read().decode()
        if p.returncode == 0:
            yield MailTest("New PGP key successfully imported:", details=gpg_output)
            yield False
        else:
            yield MailTest("An error occurs during PGP key importation:", details=gpg_output)


def assume_rfc3156(msg):
    if msg.get_param("protocol") is None or msg.get_param("protocol") != "application/pgp-signature" or msg.get_payload(1).get_content_type() != "application/pgp-signature":
        yield MailTest("Message treated as RFC3156 due to Content-Type, but is not compliant", 1)
        return

    # Extracting signature
    try:
        data = msg.get_payload(0)
        sign = msg.get_payload(1).get_payload().encode()

        # Except an exception in the two above lines if one part doesn't exist
        yield MailTest("Message treated as RFC3156: content and signature found")

        yield (data, sign)

    except IndexError:
        yield MailTest("Message treated as RFC3156 due to Content-Type, but is not compliant", 1)
        return


def assume_oldstyle(payload):
    yield MailTest("Found BEGIN PGP SIGNED MESSAGE: message treated as old style PGP email.")
    yield payload



def check(msg, GNUPG_DIRECTORY, beta=False):
    ct = msg.get_content_type()

    # First, looking for public key
    for part in msg.walk():
        if part.get_content_type() == "application/pgp-keys" and not part.is_multipart() and part.get_payload(decode=True).find(b"-----BEGIN PGP PUBLIC KEY BLOCK-----") >= 0:
            if part.get_content_type() != "application/pgp-keys":
                yield MailTest("Public key file discovered, but content-type mismatched: got %s instead of application/pgp-keys." % part.get_content_type(), 2)
            yield from import_pubkey(part.get_payload(decode=True), GNUPG_DIRECTORY)
            return

    if ct == "multipart/signed" and msg.is_multipart():
        yield from assume_rfc3156(msg)

    else:
        yield MailTest("This is not a signed e-mail: %s." % ct, 1)

        if ct == "multipart/encrypted":
            yield MailTest("As an automated service, I can't access my owner's private key. Please resend your email, unencrypted but signed.", -1)
            return

        from archive import _guess_mime

        # Looking for signed content
        for part in msg.walk():
            payload = part.get_payload()
            if part.get_content_type() == "application/pgp-signature" or (
                    payload is not None and not part.is_multipart() and part.get_payload(decode=True).find(b"-----BEGIN PGP SIGNATURE-----") >= 0
            ):
                if part.get_content_type() != "application/pgp-signature":
                    yield MailTest("Standalone PGP signature file discovered, but content-type mismatched: got %s instead of application/pgp-signature." % part.get_content_type(), 2)
                p = [x for x in msg.walk()]
                for s in range(len(p) - 1, -1, -1):
                    spart = p[s]
                    if part is not spart and not spart.is_multipart():
                        yield MailTest("Separate signature found. Trying it with part %d (%s) ..." % (s, spart.get_content_type()), -1)
                        yield (spart.get_payload(decode=True), part.get_payload(decode=True))

            elif payload is not None and not part.is_multipart() and part.get_payload(decode=True).find(b"-----BEGIN PGP PUBLIC KEY BLOCK-----") >= 0:
                if part.get_content_type() != "application/pgp-keys":
                    yield MailTest("Public key file discovered, but content-type mismatched: got %s instead of application/pgp-keys." % part.get_content_type(), 2)
                yield from import_pubkey(part.get_payload(decode=True), GNUPG_DIRECTORY)
                return

            elif payload is not None and not part.is_multipart() and part.get_payload(decode=True).find(b"-----BEGIN PGP SIGNED MESSAGE-----") >= 0:
                res = re.match(".*(-----BEGIN PGP SIGNED MESSAGE-----(.*)-----BEGIN PGP SIGNATURE-----(.*)-----END PGP SIGNATURE-----).*", payload, re.DOTALL)
                if res is not None:
                    yield from assume_oldstyle(payload)
                else:
                    res = re.match(b".*(-----BEGIN PGP SIGNED MESSAGE-----(.*)-----BEGIN PGP SIGNATURE-----(.*)-----END PGP SIGNATURE-----).*", part.get_payload(decode=True), re.DOTALL)
                    if res is not None:
                        yield from assume_oldstyle(part.get_payload(decode=True))