Source code for ferrosoft.apps.befundung.sftp

import base64
from contextlib import contextmanager
import io
import tempfile
from typing import Iterator, Tuple

import paramiko

from ferrosoft.apps.befundung.models import Report, Settings




def load_key(conf: "Settings") -> Tuple[paramiko.RSAKey, paramiko.HostKeys]:
    if conf.sftp_client_key == "":
        raise ValueError("Private key is empty")

    if conf.sftp_host_key == "":
        raise ValueError("Host key is empty")

    # Load RSA client key from in-memory buffer (decrypt if password is supplied).
    client_key = paramiko.RSAKey.from_private_key(
        io.StringIO(conf.sftp_client_key),
        password=(conf.sftp_client_key_password or None),
    )

    # Load pinned RSA host key from base64 buffer.
    pinned_host_key = paramiko.RSAKey(data=base64.b64decode(conf.sftp_host_key))
    host_keys = paramiko.HostKeys()

    # RSA is hard coded, because it was a restriction with the previously
    # used package. Supporting other key types would require model changes,
    # which is avoided for now.
    host_keys.add(conf.sftp_hostname, "ssh-rsa", pinned_host_key)

    # Non-standard port requires special host key entry.
    if getattr(conf, "sftp_port", 22) not in (None, 22):
        host_keys.add(
            f"[{conf.sftp_hostname}]:{conf.sftp_port}", "ssh-rsa", pinned_host_key
        )

    return client_key, host_keys





@contextmanager
def sftp_connection(
    conf: Settings, *, timeout: float = 30.0
) -> Iterator[paramiko.SFTPClient]:
    client_key, pinned_hostkeys = load_key(conf)

    ssh = paramiko.SSHClient()

    # Strict pinning: trust ONLY the pinned keys (do not load system/user known_hosts).
    ssh.get_host_keys().clear()
    for host, keydict in pinned_hostkeys.items():
        for keytype, key in keydict.items():
            ssh.get_host_keys().add(host, keytype, key)

    ssh.set_missing_host_key_policy(paramiko.RejectPolicy())

    try:
        ssh.connect(
            hostname=conf.sftp_hostname,
            port=conf.sftp_port,
            username=conf.sftp_username,
            pkey=client_key,
            look_for_keys=False,
            allow_agent=False,
            timeout=timeout,
        )
        with ssh.open_sftp() as sftp:
            yield sftp
    finally:
        ssh.close()





def put_examination_report(report: Report, conf: Settings, pdf_blob: bytes):
    if not conf.sftp_enable_document_upload and not conf.sftp_enable_image_upload:
        return

    if conf.sftp_hostname == "":
        raise ValueError("SFTP hostname is empty")

    with sftp_connection(conf) as sftp:
        if conf.sftp_enable_document_upload:
            with tempfile.NamedTemporaryFile() as pdf_file:
                file_name = conf.format_report_pdf_name(report)
                # Write pdf contents to temporary file, because pysftp only supports normal files.
                pdf_file.write(pdf_blob)
                pdf_file.seek(0)
                sftp.put(
                    pdf_file.name,
                    "%s/%s" % (conf.sftp_remote_dir, file_name),
                )

        if conf.sftp_enable_image_upload:
            for i, image in enumerate(report.images):
                image_name = conf.format_image_name(report, i, image.media_type)
                with tempfile.NamedTemporaryFile() as image_file:
                    image_file.write(image.binary_data)
                    image_file.seek(0)
                    sftp.put(
                        image_file.name,
                        "%s/%s" % (conf.sftp_remote_dir, image_name),
                    )