paieperless/main.py

import argparse
import imaplib
from dataclasses import dataclass
from email.policy import default
from email.parser import BytesParser
import os
import gnupg
import requests
from urllib.parse import urljoin


class Script:
    def __init__(
        self,
        imap_server,
        email_account,
        email_password,
        pgp_passphrase,
        paperless_ngx_url,
        paperless_api_token,
        gpg_key_file,
        gpg_content,
        gnupg_home,
    ):
        self.gpg = gnupg.GPG(gnupghome=gnupg_home)
        self.imap_server = imap_server
        self.email_account = email_account
        self.email_password = email_password
        self.pgp_passphrase = pgp_passphrase
        self.paperless_ngx_url = paperless_ngx_url
        self.paperless_api_token = paperless_api_token

        if gpg_key_file:
            self.gpg.import_keys_file(gpg_key_file)
        elif gpg_content:
            import_result = self.gpg.import_keys(gpg_content)
            print(f"Imported keys: {import_result.count}")
            print(f"Fingerprints: {import_result.fingerprints}")

    def decrypt_file(self, filename, filepath):
        decrypted_data = self.gpg.decrypt_file(
            open(filepath, "rb"), passphrase=self.pgp_passphrase
        )

        if not decrypted_data.ok:
            print("Decryption failed")
            return False, False

        decrypted_filename = (
            f"{os.path.splitext(filename)[0]}_decrypted{os.path.splitext(filename)[1]}"
        )
        decrypted_filepath = os.path.join("/tmp", decrypted_filename)

        with open(decrypted_filepath, "wb") as f:
            f.write(bytes(decrypted_data.data))

        return decrypted_filename, decrypted_filepath

    def upload_to_paperless(self, filename, filepath):
        with open(filepath, "rb") as f:
            files = {"document": (filename, f)}
            headers = {
                "Authorization": f"Token {self.paperless_api_token}",
            }

            response = requests.post(
                urljoin(self.paperless_ngx_url, "/api/documents/post_document/"),
                files=files,
                headers=headers,
            )

            if response.ok:
                print(f"Uploaded {filename} to Paperless-ngx")
            else:
                print(f"Failed to upload {filename}: {response.text}")

    def process_part(self, part):
        if part.get_content_maintype() == "multipart":
            return
        if part.get("Content-Disposition") is None:
            return

        filename = part.get_filename()
        if not filename:
            return

        filepath = os.path.join("/tmp", filename)
        with open(filepath, "wb") as f:
            f.write(part.get_payload(decode=True))

        decrypted_filename, decrypted_filepath = self.decrypt_file(filename, filepath)
        if not decrypted_filename:
            return

        self.upload_to_paperless(decrypted_filename, decrypted_filepath)

        os.remove(decrypted_filepath)

    def read_messages(self):
        mail = imaplib.IMAP4_SSL(self.imap_server)
        mail.login(self.email_account, self.email_password)
        mail.select("inbox")

        status, messages = mail.search(None, '(UNSEEN SUBJECT "Paycheck")')
        messages = messages[0].split()
        for msg_id in messages:
            status, data = mail.fetch(msg_id, "(RFC822)")
            raw_email = data[0][1]
            msg = BytesParser(policy=default).parsebytes(raw_email)

            for part in msg.walk():
                self.process_part(part)

        mail.logout()


@dataclass
class Argument:
    terminal_arg_name: str
    environ_arg_name: str
    help_string: str
    default: any = None
    required: bool = True

    def get_value(self, parser_args):
        arg = parser_args.get(self.terminal_arg_name, None)
        if arg:
            return arg

        arg = os.environ.get(self.environ_arg_name, None)
        if arg:
            return arg

        if self.required:
            raise Exception(
                f"Missing parameter {self.terminal_arg_name} or environment variable {self.environ_arg_name}"
            )

        return self.default


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    args = {
        "imap_server": Argument(
            terminal_arg_name="--imap-server",
            environ_arg_name="IMAP_SERVER",
            help_string="IMAP server address",
        ),
        "email_account": Argument(
            terminal_arg_name="--email-account",
            environ_arg_name="EMAIL_ACCOUNT",
            help_string="Email address",
        ),
        "email_password": Argument(
            terminal_arg_name="--email-password",
            environ_arg_name="EMAIL_PASSWORD",
            help_string="Email password",
        ),
        "pgp_passphrase": Argument(
            terminal_arg_name="--pgp-passphrase",
            environ_arg_name="PGP_PASSPHRASE",
            help_string="PGP passphrase",
        ),
        "paperless_ngx_url": Argument(
            terminal_arg_name="--paperless-ngx-url",
            environ_arg_name="PAPERLESS_NGX_URL",
            help_string="Paperless-ngx URL",
        ),
        "paperless_api_token": Argument(
            terminal_arg_name="--paperless-api-token",
            environ_arg_name="PAPERLESS_API_TOKEN",
            help_string="Paperless-ngx API token",
        ),
        "gpg_key_file": Argument(
            terminal_arg_name="--gpg-key-file",
            environ_arg_name="GPG_KEY_FILE",
            help_string="GPG key file",
            required=False,
        ),
        "gpg_content": Argument(
            terminal_arg_name="--gpg-content",
            environ_arg_name="GPG_CONTENT",
            help_string="GPG key file content",
            required=False,
        ),
        "gnupg_home": Argument(
            terminal_arg_name="--gnupg-home",
            environ_arg_name="GNUPG_HOME",
            help_string="GnuPG home directory",
        ),
    }

    for _, argument in args.items():
        parser.add_argument(argument.terminal_arg_name, help=argument.help_string)

    parser_args = vars(parser.parse_args())

    script = Script(
        **{key: argument.get_value(parser_args) for key, argument in args.items()}
    )

    script.read_messages()