# Copyright (c) 2016 Thomas Nicholson <tnnich@googlemail.com>
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. The names of the author(s) may not be used to endorse or promote
#    products derived from this software without specific prior written
#    permission.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHORS ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.

from __future__ import annotations

import hashlib
import os

from twisted.python import log
from twisted.conch.ssh import filetransfer

from cowrie.core.config import CowrieConfig
from cowrie.ssh_proxy.protocols import base_protocol


# PACKETLAYOUT = {
#     1: "SSH_FXP_INIT",
#     # ['uint32', 'version'], [['string', 'extension_name'], ['string', 'extension_data']]]
#     2: "SSH_FXP_VERSION",
#     # [['uint32', 'version'], [['string', 'extension_name'], ['string', 'extension_data']]]
#     3: "SSH_FXP_OPEN",
#     # [['uint32', 'id'], ['string', 'filename'], ['uint32', 'pflags'], ['ATTRS', 'attrs']]
#     4: "SSH_FXP_CLOSE",  # [['uint32', 'id'], ['string', 'handle']]
#     5: "SSH_FXP_READ",  # [['uint32', 'id'], ['string', 'handle'], ['uint64', 'offset'], ['uint32', 'len']]
#     6: "SSH_FXP_WRITE",
#     # [['uint32', 'id'], ['string', 'handle'], ['uint64', 'offset'], ['string', 'data']]
#     7: "SSH_FXP_LSTAT",  # [['uint32', 'id'], ['string', 'path']]
#     8: "SSH_FXP_FSTAT",  # [['uint32', 'id'], ['string', 'handle']]
#     9: "SSH_FXP_SETSTAT",  # [['uint32', 'id'], ['string', 'path'], ['ATTRS', 'attrs']]
#     10: "SSH_FXP_FSETSTAT",  # [['uint32', 'id'], ['string', 'handle'], ['ATTRS', 'attrs']]
#     11: "SSH_FXP_OPENDIR",  # [['uint32', 'id'], ['string', 'path']]
#     12: "SSH_FXP_READDIR",  # [['uint32', 'id'], ['string', 'handle']]
#     13: "SSH_FXP_REMOVE",  # [['uint32', 'id'], ['string', 'filename']]
#     14: "SSH_FXP_MKDIR",  # [['uint32', 'id'], ['string', 'path'], ['ATTRS', 'attrs']]
#     15: "SSH_FXP_RMDIR",  # [['uint32', 'id'], ['string', 'path']]
#     16: "SSH_FXP_REALPATH",  # [['uint32', 'id'], ['string', 'path']]
#     17: "SSH_FXP_STAT",  # [['uint32', 'id'], ['string', 'path']]
#     18: "SSH_FXP_RENAME",  # [['uint32', 'id'], ['string', 'oldpath'], ['string', 'newpath']]
#     19: "SSH_FXP_READLINK",  # [['uint32', 'id'], ['string', 'path']]
#     20: "SSH_FXP_SYMLINK",  # [['uint32', 'id'], ['string', 'linkpath'], ['string', 'targetpath']]
#     101: "SSH_FXP_STATUS",
#     # [['uint32', 'id'], ['uint32', 'error_code'], ['string', 'error_message'], ['string', 'language']]
#     102: "SSH_FXP_HANDLE",  # [['uint32', 'id'], ['string', 'handle']]
#     103: "SSH_FXP_DATA",  # [['uint32', 'id'], ['string', 'data']]
#     104: "SSH_FXP_NAME",
#     # [['uint32', 'id'], ['uint32', 'count'], [['string', 'filename'], ['string', 'longname'], ['ATTRS', 'attrs']]]
#     105: "SSH_FXP_ATTRS",  # [['uint32', 'id'], ['ATTRS', 'attrs']]
#     200: "SSH_FXP_EXTENDED",  # []
#     201: "SSH_FXP_EXTENDED_REPLY",  # []
# }


class SFTP(base_protocol.BaseProtocol):
    prevID: int = 0
    ID: int = 0
    handle: bytes = b""
    path: bytes = b""
    command: bytes = b""
    payloadSize: int = 0
    payloadOffset: int = 0
    theFile: bytes = b""
    parentPacket: base_protocol.BaseProtocol

    def __init__(self, uuid, chan_name, ssh):
        super().__init__(uuid, chan_name, ssh)

        self.downloadPath: str = CowrieConfig.get("honeypot", "download_path")

        self.clientPacket = base_protocol.BaseProtocol()
        self.serverPacket = base_protocol.BaseProtocol()

        self.parent: str
        self.offset: int = 0

    def parse_packet(self, parent: str, data: bytes) -> None:
        self.parent = parent

        if parent == "[SERVER]":
            self.parentPacket = self.serverPacket
        elif parent == "[CLIENT]":
            self.parentPacket = self.clientPacket
        else:
            raise ValueError

        if self.parentPacket.packetSize == 0:
            self.parentPacket.packetSize = int(data[:4].hex(), 16) - len(data[4:])
            data = data[4:]
            self.parentPacket.data = data
            data = b""

        else:
            if len(data) > self.parentPacket.packetSize:
                self.parentPacket.data = (
                    self.parentPacket.data + data[: self.parentPacket.packetSize]
                )
                data = data[self.parentPacket.packetSize :]
                self.parentPacket.packetSize = 0
            else:
                self.parentPacket.packetSize -= len(data)
                self.parentPacket.data = self.parentPacket.data + data
                data = b""

        if self.parentPacket.packetSize == 0:
            self.handle_packet(parent)

        if len(data) != 0:
            self.parse_packet(parent, data)

    def handle_packet(self, parent: str) -> None:
        self.packetSize: int = self.parentPacket.packetSize
        self.data: bytes = self.parentPacket.data
        self.command: bytes

        sftp_num: int = self.extract_int(1)

        self.prevID: int = self.ID
        self.ID: int = self.extract_int(4)

        self.path: bytes = b""

        if sftp_num == filetransfer.FXP_OPENDIR:
            self.path = self.extract_string()

        elif sftp_num == filetransfer.FXP_REALPATH:
            self.path = self.extract_string()
            self.command = b"cd " + self.path
            log.msg(parent + "[SFTP] Entered Command: " + self.command.decode())

        elif sftp_num == filetransfer.FXP_OPEN:
            self.path = self.extract_string()
            pflags = f"{self.extract_int(4):08b}"

            if pflags[6] == "1":
                self.command = b"put " + self.path
                self.theFile = b""
                # self.out.download_started(self.uuid, self.path)
            elif pflags[7] == "1":
                self.command = b"get " + self.path
            else:
                # Unknown PFlag
                log.msg(
                    parent + f"[SFTP] New SFTP pflag detected: {pflags!r} {self.data!r}"
                )

            log.msg(parent + " [SFTP] Entered Command: " + self.command.decode())

        elif sftp_num == filetransfer.FXP_READ:
            pass

        elif sftp_num == filetransfer.FXP_WRITE:
            if self.handle == self.extract_string():
                self.offset = self.extract_int(8)
                self.theFile = self.theFile[: self.offset] + self.extract_data()

        elif sftp_num == filetransfer.FXP_HANDLE:
            if self.ID == self.prevID:
                self.handle = self.extract_string()

        elif sftp_num == filetransfer.FXP_READDIR:
            if self.handle == self.extract_string():
                self.command = b"ls " + self.path

        elif sftp_num == filetransfer.FXP_SETSTAT:
            self.path = self.extract_string()
            self.command = self.extract_attrs() + b" " + self.path

        elif sftp_num == filetransfer.FXP_EXTENDED:
            cmd = self.extract_string()
            self.path = self.extract_string()

            if cmd == b"statvfs@openssh.com":
                self.command = b"df " + self.path
            elif cmd == b"hardlink@openssh.com":
                self.command = b"ln " + self.path + b" " + self.extract_string()
            elif cmd == b"posix-rename@openssh.com":
                self.command = b"mv " + self.path + b" " + self.extract_string()
            else:
                # UNKNOWN COMMAND
                log.msg(
                    parent
                    + f"[SFTP] New SFTP Extended Command detected: {cmd!r} {self.data!r}"
                )

        elif sftp_num == filetransfer.FXP_EXTENDED_REPLY:
            log.msg(parent + " [SFTP] Entered Command: " + self.command.decode())
            # self.out.command_entered(self.uuid, self.command)

        elif sftp_num == filetransfer.FXP_CLOSE:
            if self.handle == self.extract_string():
                if b"get" in self.command:
                    log.msg(
                        parent + " [SFTP] Finished Downloading: " + self.path.decode()
                    )
                elif b"put" in self.command:
                    log.msg(
                        parent + " [SFTP] Finished Uploading: " + self.path.decode()
                    )

                    # TODO: should use artifact functions
                    shasum = hashlib.sha256(self.theFile).hexdigest()
                    outfile = os.path.join(self.downloadPath, shasum)
                    fname = self.command.decode().split(" ")[-1]
                    duplicate = os.path.exists(outfile)

                    if not duplicate:
                        f = open(outfile, "wb")
                        f.write(self.theFile)
                        f.close()

                    log.msg(
                        format='SFTP Uploaded file "%(filename)s" to %(outfile)s',
                        eventid="cowrie.session.file_upload",
                        filename=fname,
                        duplicate=duplicate,
                        url=fname,
                        outfile=outfile,
                        shasum=shasum,
                        destfile=fname,
                    )

                    # if self.out.cfg.getboolean(['download', 'passive']):
                    #     # self.out.make_downloads_folder()
                    #     outfile = self.out.downloadFolder + datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f")\
                    #     + "-" + self.path.split('/')[-1]
                    #     f = open(outfile, 'wb')
                    #     f.write(self.theFile)
                    #     f.close()
                    #     #self.out.file_downloaded((self.uuid, True, self.path, outfile, None))

        elif sftp_num == filetransfer.FXP_SYMLINK:
            self.command = (
                b"ln -s " + self.extract_string() + b" " + self.extract_string()
            )

        elif sftp_num == filetransfer.FXP_MKDIR:
            self.command = b"mkdir " + self.extract_string()

        elif sftp_num == filetransfer.FXP_REMOVE:
            self.command = b"rm " + self.extract_string()

        elif sftp_num == filetransfer.FXP_RMDIR:
            self.command = b"rmdir " + self.extract_string()

        elif sftp_num == filetransfer.FXP_STATUS:
            if self.ID == self.prevID:
                code = self.extract_int(4)
                if code in [0, 1]:
                    if b"get" not in self.command and b"put" not in self.command:
                        log.msg(
                            parent + " [SFTP] Entered Command: " + self.command.decode()
                        )
                else:
                    message = self.extract_string()
                    log.msg(
                        parent
                        + " [SFTP] Failed Command: "
                        + self.command.decode()
                        + " Reason: "
                        + message.decode()
                    )
        else:
            log.msg("[SFTP] Unhandled packet: {sftp_num}")

    def extract_attrs(self) -> bytes:
        cmd: str = ""
        flags: str = f"{self.extract_int(4):08b}"

        if flags[5] == "1":
            perms = f"{self.extract_int(4):09b}"
            # log.msg(log.LPURPLE, self.parent + '[SFTP]', 'PERMS:' + perms)
            chmod = (
                str(int(perms[:3], 2))
                + str(int(perms[3:6], 2))
                + str(int(perms[6:], 2))
            )
            cmd = "chmod " + chmod
        elif flags[6] == "1":
            user = str(self.extract_int(4))
            group = str(self.extract_int(4))
            cmd = "chown " + user + ":" + group
        else:
            pass
            # Unknown attribute
            # log.msg(log.LRED, self.parent + '[SFTP]',
            #         'New SFTP Attribute detected - Please raise a HonSSH issue on github with the details: %s %s' %
            #         (flags, self.data))
        return cmd.encode()


"""
CLIENT                              SERVER

    SSH_FXP_INIT    -->
                    <--    SSH_FXP_VERSION

    SSH_FXP_OPEN    -->
                    <--    SSH_FXP_HANDLE (or SSH_FXP_STATUS if fail)

    SSH_FXP_READ    -->
                    <--    SSH_FXP_DATA (or SSH_FXP_STATUS if fail)

    SSH_FXP_WRITE   -->
                    <--    SSH_FXP_STATUS

    SSH_FXP_REMOVE  -->
                    <--    SSH_FXP_STATUS

    SSH_FXP_RENAME  -->
                    <--    SSH_FXP_STATUS

    SSH_FXP_MKDIR   -->
                    <--    SSH_FXP_STATUS

    SSH_FXP_RMDIR   -->
                    <--    SSH_FXP_STATUS

    SSH_FXP_OPENDIR -->
                    <--    SSH_FXP_HANDLE (or SSH_FXP_STATUS if fail)

    SSH_FXP_READDIR -->
                    <--    SSH_FXP_NAME (or SSH_FXP_STATUS if fail)

    SSH_FXP_STAT    -->         //Follows symlinks
                    <--    SSH_FXP_ATTRS (or SSH_FXP_STATUS if fail)

    SSH_FXP_LSTAT    -->         //Does not follow symlinks
                    <--    SSH_FXP_ATTRS (or SSH_FXP_STATUS if fail)

    SSH_FXP_FSTAT    -->         //Works on an open file/handle not a file path like (L)STAT
                    <--    SSH_FXP_ATTRS (or SSH_FXP_STATUS if fail)

    SSH_FXP_SETSTAT -->          //Sets file attributes on path
                    <--    SSH_FXP_STATUS

    SSH_FXP_FSETSTAT-->          //Sets file attributes on a handle
                    <--    SSH_FXP_STATUS

    SSH_FXP_READLINK -->        //Used to find the target of a symlink
                    <--    SSH_FXP_NAME (or SSH_FXP_STATUS if fail)

    SSH_FXP_SYMLINK  -->        //Used to create a symlink
                    <--    SSH_FXP_NAME (or SSH_FXP_STATUS if fail)

    SSH_FXP_REALPATH -->          //Relative path
                    <--    SSH_FXP_NAME (or SSH_FXP_STATUS if fail)

    SSH_FXP_CLOSE   -->                     //Closes handle not session
                    <--    SSH_FXP_STATUS
"""
