simplified function by using map()

[edward.git] / edward-bot

#! /usr/bin/env python3

"""*********************************************************************
* 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+)  *
*                                                                      *
* Special thanks to Josh Drake for writing the original edward bot! :) *
*                                                                      *
************************************************************************

Code used from:

* http://agpl.fsf.org/emailselfdefense.fsf.org/edward/CURRENT/edward.tar.gz

"""

import sys
import email.parser
import gpgme
import re
import io
import time


def main ():

    handle_args()

    email_text = sys.stdin.read()

    email_from, email_subject = email_from_subject(email_text)

    plaintext, keys = email_decode_flatten (email_text)
    encrypt_to_key = choose_reply_encryption_key(keys)

    reply_message = generate_reply(plaintext, email_from, \
                                   email_subject, encrypt_to_key,
                                   "DAB4F989E2788B8DF058E0EFEF1EC52039B36E58")

    print(reply_message)


def email_decode_flatten (email_text):

    body = ""
    keys = []

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

    for subpart in email_struct.walk():

        payload, description, filename, content_type \
                = get_email_subpart_info(subpart)

        if payload == "":
            continue

        if content_type == "multipart":
            continue

        if content_type == "application/pgp-encrypted":
            if description == "PGP/MIME version identification":
                if payload.strip() != "Version: 1":
                    print(progname + ": Warning: unknown " \
                            + description + ": " \
                            + payload.strip(), file=sys.stderr)
            continue


        if (filename == "encrypted.asc") or (content_type == "pgp/mime"):
            plaintext, more_keys = decrypt_text(payload)

            body += plaintext
            keys += more_keys

        elif content_type == "text/plain":
            body += payload + "\n"

        else:
            body += payload + "\n"

    return body, keys


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 get_email_subpart_info (part):

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

    filename            = part.get_filename()
    content_type        = part.get_content_type()
    description_list    = part.get_params(header='content-description')

    if charset == None:
        charset = 'utf-8'

    if payload_bytes != None:
        payload = payload_bytes.decode(charset)
    else:
        payload = ""

    if description_list != None:
        description = description_list[0][0]
    else:
        description = ""

    return payload, description, filename, content_type


def decrypt_text (gpg_text):

    body = ""
    keys = []

    gpg_chunks = split_message(gpg_text)

    plaintext_and_sigs_chunks = decrypt_chunks(gpg_chunks)

    for chunk in plaintext_and_sigs_chunks:
        plaintext   = chunk[0]
        sigs        = chunk[1]

        for sig in sigs:
            key = get_pub_key(sig)
            keys += [key]

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

        body += plaintext
        keys += more_keys

    return body, keys


def get_pub_key (sig):

    gpgme_ctx = gpgme.Context()

    fingerprint = sig.fpr
    key = gpgme_ctx.get_key(fingerprint)

    return key


def split_message (text):

    gpg_matches = re.search( \
            '(-----BEGIN PGP MESSAGE-----' + \
            '.*' + \
            '-----END PGP MESSAGE-----)', \
            text, \
            flags=re.DOTALL)

    if gpg_matches != None:
        gpg_chunks = gpg_matches.groups()
    else:
        gpg_chunks = ()

    return gpg_chunks


def decrypt_chunks (gpg_chunks):

    return map(decrypt_chunk, gpg_chunks)


def decrypt_chunk (gpg_chunk):

    gpgme_ctx = gpgme.Context()

    chunk_b = io.BytesIO(gpg_chunk.encode('ASCII'))
    plain_b = io.BytesIO()

    sigs = gpgme_ctx.decrypt_verify(chunk_b, plain_b)

    plaintext = plain_b.getvalue().decode('ASCII')
    return (plaintext, sigs)


def choose_reply_encryption_key (keys):

    reply_key = None
    for key in keys:
        if (key.can_encrypt == True):
            reply_key = key
            break

    return key


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

    plaintext_reply  = "thanks for the message!\n\n\n"
    plaintext_reply += email_quote_text(plaintext)

    encrypted_reply = encrypt_sign_message(plaintext_reply, encrypt_to_key,
                                           sign_with_fingerprint)

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

    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, sign_with_fingerprint):

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

    sign_with_key = gpgme_ctx.get_key(sign_with_fingerprint)
    gpgme_ctx.signers = [sign_with_key]

    plaintext_bytes = io.BytesIO(plaintext.encode('UTF-8'))
    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 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()