#! /usr/bin/env python3
# -*- coding: utf-8 -*-
"""*********************************************************************
* Edward is free software: you can redistribute it and/or modify       *
* it under the terms of the GNU Affero Public License as published by  *
* the Free Software Foundation, either version 3 of the License, or    *
* (at your option) any later version.                                  *
*                                                                      *
* Edward is distributed in the hope that it will be useful,            *
* but WITHOUT ANY WARRANTY; without even the implied warranty of       *
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the        *
* GNU Affero Public License for more details.                          *
*                                                                      *
* You should have received a copy of the GNU Affero Public License     *
* along with Edward.  If not, see <http://www.gnu.org/licenses/>.      *
*                                                                      *
* Copyright (C) 2014-2015 Andrew Engelbrecht                (AGPLv3+)  *
* Copyright (C) 2014      Josh Drake                        (AGPLv3+)  *
* Copyright (C) 2014      Lisa Marie Maginnis               (AGPLv3+)  *
* Copyright (C) 2009-2015 Tails developers <tails@boum.org> ( GPLv3+)  *
* Copyright (C) 2009      W. Trevor King <wking@drexel.edu> ( GPLv2+)  *
*                                                                      *
* Special thanks to Josh Drake for writing the original edward bot! :) *
*                                                                      *
************************************************************************

Code sourced from these projects:

  * http://agpl.fsf.org/emailselfdefense.fsf.org/edward/CURRENT/edward.tar.gz
  * https://git-tails.immerda.ch/whisperback/tree/whisperBack/encryption.py?h=feature/python3
  * http://www.physics.drexel.edu/~wking/code/python/send_pgp_mime
"""

import sys
import gpgme
import re
import io
import os

import email.parser
import email.message
import email.encoders

from email.mime.multipart       import MIMEMultipart
from email.mime.application     import MIMEApplication
from email.mime.nonmultipart    import MIMENonMultipart

import edward_config

match_types =  [('clearsign',
                '-----BEGIN PGP SIGNED MESSAGE-----.*?-----BEGIN PGP SIGNATURE-----.*?-----END PGP SIGNATURE-----'),
                ('message',
                '-----BEGIN PGP MESSAGE-----.*?-----END PGP MESSAGE-----'),
                ('pubkey',
                '-----BEGIN PGP PUBLIC KEY BLOCK-----.*?-----END PGP PUBLIC KEY BLOCK-----'),
                ('detachedsig',
                '-----BEGIN PGP SIGNATURE-----.*?-----END PGP SIGNATURE-----')]


class EddyMsg (object):
    def __init__(self):
        self.multipart          = False
        self.subparts           = []

        self.charset            = None
        self.payload_bytes      = None
        self.payload_pieces     = []

        self.filename           = None
        self.content_type       = None
        self.description_list   = None


class PayloadPiece (object):
    def __init__(self):
        self.piece_type         = None
        self.string             = None
        self.gpg_data           = None


class GPGData (object):
    def __init__(self):
        self.decrypted          = False

        self.plainobj           = None
        self.sigs               = []
        self.keys               = []


def main ():

    handle_args()

    gpgme_ctx = get_gpg_context(edward_config.gnupghome,
                              edward_config.sign_with_key)

    email_text = sys.stdin.read()
    result = parse_pgp_mime(email_text, gpgme_ctx)

    email_from, email_subject = email_from_subject(email_text)

#    plaintext, fingerprints = email_decode_flatten(email_text, gpgme_ctx, False)
#    encrypt_to_key = choose_reply_encryption_key(gpgme_ctx, fingerprints)
#
#    reply_message = generate_reply(plaintext, email_from, \
#                                   email_subject, encrypt_to_key,
#                                   gpgme_ctx)

    print(flatten_eddy(result))


def get_gpg_context (gnupghome, sign_with_key_fp):

    os.environ['GNUPGHOME'] = gnupghome

    gpgme_ctx = gpgme.Context()
    gpgme_ctx.armor = True

    try:
        sign_with_key = gpgme_ctx.get_key(sign_with_key_fp)
    except:
        error("unable to load signing key. is the gnupghome "
                + "and signing key properly set in the edward_config.py?")
        exit(1)

    gpgme_ctx.signers = [sign_with_key]

    return gpgme_ctx


def parse_pgp_mime (email_text, gpgme_ctx):

    email_struct = email.parser.Parser().parsestr(email_text)

    eddy_obj = parse_mime(email_struct)
    eddy_obj = split_payloads(eddy_obj)
    eddy_obj = decrypt_payloads(eddy_obj, gpgme_ctx)

    return eddy_obj


def parse_mime(msg_struct):

    eddy_obj = EddyMsg()

    if msg_struct.is_multipart() == True:
        payloads = msg_struct.get_payload()

        eddy_obj.multipart = True
        eddy_obj.subparts = list(map(parse_mime, payloads))

    else:
        eddy_obj = get_subpart_data(msg_struct)

    return eddy_obj


def split_payloads (eddy_obj):

    if eddy_obj.multipart == True:
        eddy_obj.subparts = list(map(split_payloads, eddy_obj.subparts))

    else:
        for (match_type, pattern) in match_types:

            new_pieces_list = []
            for payload_piece in eddy_obj.payload_pieces:
                new_pieces_list += scan_and_split(payload_piece,
                    match_type, pattern)
            eddy_obj.payload_pieces = new_pieces_list

    return eddy_obj


def scan_and_split (payload_piece, match_type, pattern):

    flags = re.DOTALL | re.MULTILINE
    matches = re.search("(?P<beginning>.*?)(?P<match>" + pattern +
                        ")(?P<rest>.*)", payload_piece.string, flags=flags)

    if matches == None:
        pieces = [payload_piece]

    else:

        beginning               = PayloadPiece()
        beginning.string        = matches.group('beginning')
        beginning.piece_type    = payload_piece.piece_type

        match                   = PayloadPiece()
        match.string            = matches.group('match')
        match.piece_type        = match_type

        rest                    = PayloadPiece()
        rest.string             = matches.group('rest')
        rest.piece_type         = payload_piece.piece_type

        more_pieces = scan_and_split(rest, match_type, pattern)
        pieces = [beginning, match ] + more_pieces

    return pieces


def get_subpart_data (part):

    obj = EddyMsg()

    obj.charset             = part.get_content_charset()
    obj.payload_bytes       = part.get_payload(decode=True)

    obj.filename            = part.get_filename()
    obj.content_type        = part.get_content_type()
    obj.description_list    = part['content-description']

    # your guess is as good as a-myy-ee-ine...
    if obj.charset == None:
        obj.charset = 'utf-8'

    if obj.payload_bytes != None:
        try:
            payload = PayloadPiece()
            payload.string = obj.payload_bytes.decode(obj.charset)
            payload.piece_type = 'text'

            obj.payload_pieces = [payload]
        except UnicodeDecodeError:
            pass

    return obj


def do_to_eddys_pieces (function_to_do, eddy_obj, data):

    if eddy_obj.multipart == True:
        result_list = []
        for sub in eddy_obj.subparts:
            result_list += do_to_eddys_pieces(function_to_do, sub, data)
    else:
        result_list = [function_to_do(eddy_obj.payload_pieces, data)]

    return result_list


def decrypt_payloads (eddy_obj, gpgme_ctx):

    do_to_eddys_pieces(decrypt_payload_pieces, eddy_obj, gpgme_ctx)

    return eddy_obj


def decrypt_payload_pieces (payload_pieces, gpgme_ctx):

    for piece in payload_pieces:

        if piece.piece_type == "text":
            # don't transform the plaintext.
            pass

        elif piece.piece_type == "message":
            (plaintext, sigs) = decrypt_block (piece.string, gpgme_ctx)

            if plaintext:
                piece.gpg_data = GPGData()
                piece.gpg_data.sigs = sigs
                # recurse!
                piece.gpg_data.plainobj = parse_pgp_mime(plaintext, gpgme_ctx)
        else:
            pass


def flatten_eddy (eddy_obj):

    string = "\n".join(do_to_eddys_pieces(flatten_payload_pieces, eddy_obj, None))

    return string


def flatten_payload_pieces (payload_pieces, _ignore):

    string = ""
    for piece in payload_pieces:
        if piece.piece_type == "text":
            string += piece.string
        elif piece.piece_type == "message":
            # recursive!
            string += flatten_eddy(piece.gpg_data.plainobj)

    return string


def email_from_subject (email_text):

    email_struct = email.parser.Parser().parsestr(email_text)

    email_from      = email_struct['From']
    email_subject   = email_struct['Subject']

    return email_from, email_subject


def add_gpg_keys (text, gpgme_ctx):

    key_blocks = scan_and_grab(text,
                               '-----BEGIN PGP PUBLIC KEY BLOCK-----',
                               '-----END PGP PUBLIC KEY BLOCK-----')

    fingerprints = []
    for key_block in key_blocks:
        fp = io.BytesIO(key_block.encode('ascii'))

        result = gpgme_ctx.import_(fp)
        imports = result.imports

        if imports != []:
            fingerprint = imports[0][0]
            fingerprints += [fingerprint]

            debug("added gpg key: " + fingerprint)

    return fingerprints


def decrypt_text (gpg_text, gpgme_ctx):

    body = ""
    fingerprints = []

    msg_blocks = scan_and_grab(gpg_text,
                               '-----BEGIN PGP MESSAGE-----',
                               '-----END PGP MESSAGE-----')

    plaintexts_and_sigs = decrypt_blocks(msg_blocks, gpgme_ctx)

    for pair in plaintexts_and_sigs:
        plaintext   = pair[0]
        sigs        = pair[1]

        for sig in sigs:
            fingerprints += [sig.fpr]

        # recursive for nested layers of mime and/or gpg
        plaintext, more_fps = email_decode_flatten(plaintext, gpgme_ctx, True)

        body += plaintext
        fingerprints += more_fps

    return body, fingerprints


def verify_clear_signature (text, gpgme_ctx):

    sig_blocks = scan_and_grab(text,
                               '-----BEGIN PGP SIGNED MESSAGE-----',
                               '-----END PGP SIGNATURE-----')

    fingerprints = []
    plaintext = ""

    for sig_block in sig_blocks:
        msg_fp = io.BytesIO(sig_block.encode('utf-8'))
        ptxt_fp = io.BytesIO()

        result = gpgme_ctx.verify(msg_fp, None, ptxt_fp)

        plaintext += ptxt_fp.getvalue().decode('utf-8')
        fingerprint = result[0].fpr

        fingerprints += [fingerprint]

    return plaintext, fingerprints


def scan_and_grab (text, start_text, end_text):

    matches = re.search('(' + start_text + '.*' + end_text + ')',
                        text, flags=re.DOTALL)

    if matches != None:
        match_tuple = matches.groups()
    else:
        match_tuple = ()

    return match_tuple


def decrypt_blocks (msg_blocks, gpgme_ctx):

    return [decrypt_block(block, gpgme_ctx) for block in msg_blocks]


def decrypt_block (msg_block, gpgme_ctx):

    block_b = io.BytesIO(msg_block.encode('ascii'))
    plain_b = io.BytesIO()

    try:
        sigs = gpgme_ctx.decrypt_verify(block_b, plain_b)
    except:
        return ("",[])

    plaintext = plain_b.getvalue().decode('utf-8')
    return (plaintext, sigs)


def choose_reply_encryption_key (gpgme_ctx, fingerprints):

    reply_key = None
    for fp in fingerprints:
        try:
            key = gpgme_ctx.get_key(fp)

            if (key.can_encrypt == True):
                reply_key = key
                break
        except:
            continue


    return reply_key


def generate_reply (plaintext, email_from, email_subject, encrypt_to_key,
                    gpgme_ctx):


    reply  = "To: " + email_from + "\n"
    reply += "Subject: " + email_subject + "\n"

    if (encrypt_to_key != None):
        plaintext_reply  = "thanks for the message!\n\n\n"
        plaintext_reply += email_quote_text(plaintext)

        # quoted printable encoding lets most ascii characters look normal
        # before the decrypted mime message is decoded.
        char_set = email.charset.Charset("utf-8")
        char_set.body_encoding = email.charset.QP

        # MIMEText doesn't allow setting the text encoding
        # so we use MIMENonMultipart.
        plaintext_mime = MIMENonMultipart('text', 'plain')
        plaintext_mime.set_payload(plaintext_reply, charset=char_set)

        encrypted_text = encrypt_sign_message(plaintext_mime.as_string(),
                                              encrypt_to_key,
                                              gpgme_ctx)

        control_mime = MIMEApplication("Version: 1",
                                       _subtype='pgp-encrypted',
                                       _encoder=email.encoders.encode_7or8bit)
        control_mime['Content-Description'] = 'PGP/MIME version identification'
        control_mime.set_charset('us-ascii')

        encoded_mime = MIMEApplication(encrypted_text,
                                       _subtype='octet-stream; name="encrypted.asc"',
                                       _encoder=email.encoders.encode_7or8bit)
        encoded_mime['Content-Description'] = 'OpenPGP encrypted message'
        encoded_mime['Content-Disposition'] = 'inline; filename="encrypted.asc"'
        encoded_mime.set_charset('us-ascii')

        message_mime = MIMEMultipart(_subtype="encrypted", protocol="application/pgp-encrypted")
        message_mime.attach(control_mime)
        message_mime.attach(encoded_mime)
        message_mime['Content-Disposition'] = 'inline'

        reply += message_mime.as_string()

    else:
        reply += "\n"
        reply += "Sorry, i couldn't find your key.\n"
        reply += "I'll need that to encrypt a message to you."

    return reply


def email_quote_text (text):

    quoted_message = re.sub(r'^', r'> ', text, flags=re.MULTILINE)

    return quoted_message


def encrypt_sign_message (plaintext, encrypt_to_key, gpgme_ctx):

    plaintext_bytes = io.BytesIO(plaintext.encode('ascii'))
    encrypted_bytes = io.BytesIO()

    gpgme_ctx.encrypt_sign([encrypt_to_key], gpgme.ENCRYPT_ALWAYS_TRUST,
            plaintext_bytes, encrypted_bytes)

    encrypted_txt = encrypted_bytes.getvalue().decode('ascii')
    return encrypted_txt


def error (error_msg):

    sys.stderr.write(progname + ": " + str(error_msg) + "\n")


def debug (debug_msg):

    if edward_config.debug == True:
        error(debug_msg)


def handle_args ():
    if __name__ == "__main__":

        global progname
        progname = sys.argv[0]

        if len(sys.argv) > 1:
            print(progname + ": error, this program doesn't " \
                    "need any arguments.", file=sys.stderr)
            exit(1)


main()