contentmonster/classes/connection.py

import paramiko as pikuniku  # :P

from paramiko.client import SSHClient, WarningPolicy

from io import BytesIO
from pathlib import Path
from typing import Union, Optional

import errno
import stat


class Connection:
    """Class representing an SSH/SFTP connection to a Vessel
    """

    def __init__(self, vessel):
        """Initialize a new Connection to a Vessel

        Args:
            vessel (classes.vessel.Vessel): Vessel object to open connection to
        """
        self._vessel = vessel
        self._client = SSHClient()
        self._client.load_system_host_keys()
        self._client.set_missing_host_key_policy(WarningPolicy)
        self._client.connect(vessel.address, vessel.port, vessel.username,
                             vessel.password, timeout=vessel.timeout,
                             passphrase=vessel.passphrase)
        self._transport = self._client.get_transport()
        self._transport.set_keepalive(10)
        self._sftp = self._client.open_sftp()

    def _exists(self, path: Union[str, Path]) -> bool:
        """Check if a path exists on the Vessel. Symlinks are not followed.

        Args:
            path (str, pathlib.Path): Path to check on the vessel

        Returns:
            bool: True if the path exists on the Vessel, else False
        """
        try:
            self._sftp.stat(str(path))
            return True
        except FileNotFoundError:
            return False

    def _isdir(self, path: Union[str, Path]) -> bool:
        """Check if a path is a directory on the Vessel. Symlinks are followed.

        Args:
            path (str, pathlib.Path): Path to check on the vessel

        Returns:
            bool: True if the path provided is a directory on the Vessel, False
              if it is a different kind of file.

        Raises:
            FileNotFoundError: Raised if the path does not exist on the Vessel
        """
        return stat.S_ISDIR(self._sftp.lstat(str(path)).st_mode)

    def _mkdir(self, path: Union[str, Path]) -> None:
        """Create new directory on the Vessel

        Args:
            path (str, pathlib.Path): Path at which to create a new directory
              on the Vessel

        Raises:
            IOError: Raised if the directory cannot be created
        """
        self._sftp.mkdir(str(path))

    def _listdir(self, path: Optional[Union[str, Path]] = None) -> list[Optional[str]]:
        """List files in a directory on the Vessel

        Args:
            path (str, pathlib.Path, optional): Path at which to list contents.
              Will use current working directory if None. Defaults to None.

        Returns:
            list: List of the names of files (str) located at the provided path
        """
        return self._sftp.listdir(str(path) if path else ".")

    def _remove(self, path: Union[str, Path]) -> None:
        """Remove a file from the Vessel

        Args:
            path (str, pathlib.Path): Path of the file to delete

        Raises:
            FileNotFoundError: Raised if no file is found at the given path
            IOError: Raised if the file cannot be deleted
        """
        return self._sftp.remove(str(path))

    def assertTempDirectory(self) -> None:
        """Make sure that the temp directory exists on the Vessel

        Raises:
            ValueError: Raised if the path is already in use on the vessel but
              is not a directory.
            IOError: Raised if the directory does not exist but cannot be 
              created.
        """
        if not self._exists(self._vessel.tempdir):
            self._mkdir(self._vessel.tempdir)
        elif not self._isdir(self._vessel.tempdir):
            raise ValueError(
                f"{self._vessel.tempdir} exists but is not a directory on Vessel {self._vessel.name}!")        

    def assertDirectories(self, directory) -> None:
        """Make sure that destination and temp directories exist on the Vessel

        Args:
            directory (classes.directory.Directory): Directory object
              representing the directory to check

        Raises:
            ValueError: Raised if a path is already in use on the vessel but
              not a directory.
            IOError: Raised if a directory that does not exist cannot be 
              created.
        """
        for d in [directory, self._vessel.tempdir]:
            if not self._exists(d):
                self._mkdir(d)
            elif not self._isdir(d):
                raise ValueError(
                    f"{d} exists but is not a directory on Vessel {self._vessel.name}!")

    def assertChunkComplete(self, chunk, path: Optional[Union[str, Path]] = None) -> bool:
        """Check if a Chunk has been uploaded correctly

        Args:
            chunk (classes.chunk.Chunk): Chunk object to verify upload for
            path (str, pathlib.Path, optional): Optional path at which to 
              check. If None, will get default path from Chunk object. 
              Defaults to None.

        Returns:
            bool: True if file has been uploaded correctly, else False
        """
        path = path or self._vessel.tempdir / chunk.getTempName()

        if self._exists(path):
            # "-b" should not be required, but makes sure to use binary mode
            _, o, _ = self._client.exec_command("sha256sum -b " + str(path))

            # Blocking for the command to complete
            o.channel.recv_exit_status()

            # Remove the file if it is broken
            if not o.readline().split()[0] == chunk.getHash():
                self._remove(path)

            else:
                return True
        return False

    def pushChunk(self, chunk, path: Optional[Union[str, Path]] = None) -> None:
        """Push the content of a Chunk object to the Vessel

        Args:
            chunk (classes.chunk.Chunk): Chunk object containing the data to
              push to the Vessel
            path (str, pathlib.Path, optional): Path at which to store the
              Chunk on the Vessel. If None, use default location provided by
              Chunk object. Defaults to None.
        """
        path = path or (self._vessel.tempdir / chunk.getTempName())
        flo = BytesIO(chunk.data)
        self._sftp.putfo(flo, str(path), len(chunk.data))

    def compileComplete(self, remotefile) -> None:
        """Build a complete File from uploaded Chunks.

        Args:
            remotefile (classes.remotefile.RemoteFile): RemoteFile object
              describing the uploaded File
        """
        numchunks = remotefile.getStatus() + 1

        # Get files in correct order to concatenate
        files = " ".join(
            [str(self._vessel.tempdir / f"{remotefile.file.uuid}_{i}.part") for i in range(numchunks)])

        completefile = remotefile.file.getChunk(-1)
        outname = completefile.getTempName()
        outpath = self._vessel.tempdir / outname
        _, o, _ = self._client.exec_command(f"cat {files} > {outpath}")

        # Block for command to complete
        o.channel.recv_exit_status()

    def assertComplete(self, remotefile, allow_retry: bool = False) -> bool:
        """Check if File has been reassembled from Chunks correctly

        Args:
            remotefile (classes.remotefile.RemoteFile): RemoteFile object
              describing the uploaded File
            allow_retry (bool, optional): If True, assume that compileComplete
              failed for some other reason than corrupt Chunks, and only delete
              compiled file, else clear entire temporary directory. Defaults to
              False.

        Returns:
            bool: True if file was reassembled correctly, else False
        """
        completefile = remotefile.file.getChunk(-1)
        outname = completefile.getTempName()
        outpath = self._vessel.tempdir / outname

        if not self._exists(outpath):
            return False

        if not self.assertChunkComplete(completefile):
            if allow_retry:
                self._remove(outpath)
            else:
                self.clearTempDir()
            return False

        return True

    def moveComplete(self, remotefile) -> None:
        """Moves reassembled file to output destination

        Args:
            remotefile (classes.remotefile.RemoteFile): RemoteFile object 
              describing the uploaded File.

        Returns:
            [type]: [description]
        """
        completefile = remotefile.file.getChunk(-1)
        destination = remotefile.file.getFullPath()
        self._sftp.posix_rename(
            str(self._vessel.tempdir / completefile.getTempName()), str(destination))

        # Make sure that file has actually been created at destination
        self._sftp.stat(str(destination))

    def getCurrentUploadUUID(self) -> Optional[str]:
        """Get UUID of file currently being uploaded

        Returns:
            str, optional: UUID of the File being uploaded, if any, else None
        """
        for f in self._listdir(self._vessel.tempdir):
            if f.endswith(".part"):
                return f.split("_")[0]

    def clearTempDir(self) -> None:
        """Clean up the temporary directory on the Vessel
        """
        for f in self._listdir(self._vessel.tempdir):
            self._remove(self._vessel.tempdir / f)

    def __del__(self):
        """Close SSH connection when ending Connection
        """
        self._client.close()