#!/usr/bin/env python
# -*- coding: iso-8859-15 -*-
# -*- Mode: python -*-

from __future__ import print_function

import cmd
import os
import random
import shlex
import string
import sys
import time
from subprocess import Popen, PIPE, STDOUT
from threading import Lock, Thread

from lib.common import DataStore
from lib.exceptions import keimpxError
from lib.logger import logger

try:
    from impacket.nmb import NetBIOSTimeout
    from impacket.examples import remcomsvc
    from impacket import smb
    from impacket.smbconnection import SMBConnection, SessionError
    from impacket.structure import Structure
except ImportError:
    sys.stderr.write('psexec: Impacket import error')
    sys.stderr.write('Impacket by SecureAuth Corporation is required for this tool to work. Please download it using:'
                     '\npip: pip install -r requirements.txt\nOr through your package manager:\npython-impacket.')
    sys.exit(255)


###############################################################
# Code borrowed and adapted from Impacket's psexec.py example #
###############################################################
class RemComMessage(Structure):
    structure = (
        ('Command', '4096s=""'),
        ('WorkingDir', '260s=""'),
        ('Priority', '<L=0x20'),
        ('ProcessID', '<L=0x01'),
        ('Machine', '260s=""'),
        ('NoWait', '<L=0'),
    )


class RemComResponse(Structure):
    structure = (
        ('ErrorCode', '<L=0'),
        ('ReturnCode', '<L=0'),
    )


RemComSTDOUT = "RemCom_stdout"
RemComSTDIN = "RemCom_stdin"
RemComSTDERR = "RemCom_stderr"

lock = Lock()


class PsExec(object):
    def __init__(self):
        pass

    def psexec(self, command=None):
        srvname = ''.join([random.choice(string.ascii_letters) for _ in range(8)])
        remote_file = '%s.exe' % ''.join([random.choice(string.ascii_lowercase) for _ in range(8)])

        if not command:
            logger.info('Command has not been specified, going to call cmd.exe')
            command = 'cmd.exe'

        if command in ('cmd.exe', 'command.com'):
            logger.info('Launching interactive OS shell')

        command_and_args = shlex.split(command)

        if os.path.exists(command_and_args[0]):
            self.use(DataStore.writable_share)
            self.upload(command_and_args[0])

        logger.debug('Going to use temporary service %s' % srvname)

        self.deploy(srvname, local_file=remcomsvc.RemComSvc(), srvargs='', remote_file=remote_file)
        self.smb_transport('svcctl')
        self.__smb = self.trans.get_smb_connection()
        self.__smb.setTimeout(100000)
        self.__tid = self.__smb.connectTree('IPC$')
        self.__fid_main = self.openPipe(self.__smb, self.__tid, '\\RemCom_communicaton', 0x12019f)

        packet = RemComMessage()
        packet['Machine'] = ''.join([random.choice(string.ascii_letters) for i in range(4)])
        packet['Command'] = os.path.basename(command.replace('\\', '/'))
        packet['ProcessID'] = os.getpid()

        self.__smb.writeNamedPipe(self.__tid, self.__fid_main, packet.getData())

        # Here we'll store the command we type so we don't print it back ;)
        # ( I know.. globals are nasty :P )
        global LastDataSent
        LastDataSent = ''

        # Create the pipes threads
        stdin_pipe = RemoteStdInPipe(self.trans, '\\%s%s%d' % (RemComSTDIN, packet['Machine'], packet['ProcessID']),
                                     smb.FILE_WRITE_DATA | smb.FILE_APPEND_DATA, self.share)
        stdin_pipe.start()
        stdout_pipe = RemoteStdOutPipe(self.trans, '\\%s%s%d' % (RemComSTDOUT, packet['Machine'],
                                                                 packet['ProcessID']), smb.FILE_READ_DATA)
        stdout_pipe.start()
        stderr_pipe = RemoteStdErrPipe(self.trans, '\\%s%s%d' % (RemComSTDERR, packet['Machine'],
                                                                 packet['ProcessID']), smb.FILE_READ_DATA)
        stderr_pipe.start()

        # And we stay here till the end
        ans = self.__smb.readNamedPipe(self.__tid, self.__fid_main, 8)

        if len(ans):
            retCode = RemComResponse(ans)
            logger.info('Process %s finished with ErrorCode: %d, ReturnCode: %d'
                        % (os.path.basename(command.replace('\\', '/')), retCode['ErrorCode'], retCode['ReturnCode']))

        self.undeploy(srvname)

    def openPipe(self, s, tid, pipe, accessMask):
        pipeReady = False
        tries = 50

        while pipeReady is False and tries > 0:
            try:
                self.__smb.waitNamedPipe(tid, pipe)
                pipeReady = True
            except Exception as e:
                # traceback.print_exc()
                logger.error('Named pipe open error: %s' % str(e))
                tries -= 1
                time.sleep(2)

        if tries == 0:
            logger.error('Named pipe not ready, aborting')
            raise

        fid = self.__smb.openFile(tid, pipe, accessMask, creationOption=0x40, fileAttributes=0x80)

        return fid


class Pipes(Thread):
    def __init__(self, transport, pipe, permissions, share=None):
        Thread.__init__(self)
        self.server = 0
        self.transport = transport
        self.credentials = transport.get_credentials()
        self.tid = 0
        self.fid = 0
        self.share = share
        self.port = transport.get_dport()
        self.pipe = pipe
        self.permissions = permissions
        self.daemon = True

    def connectPipe(self):
        try:
            lock.acquire()
            self.server = SMBConnection('*SMBSERVER', self.transport.get_smb_connection().getRemoteHost(),
                                        sess_port=self.port)
            user, passwd, domain, lm, nt, _, _, _ = self.credentials
            self.server.login(user, passwd, domain, lm, nt)
            lock.release()
            self.tid = self.server.connectTree('IPC$')

            self.server.waitNamedPipe(self.tid, self.pipe)
            self.fid = self.server.openFile(self.tid, self.pipe, self.permissions, creationOption=0x40,
                                            fileAttributes=0x80)
            self.server.setTimeout(1000000)
        except Exception as e:
            # traceback.print_exc()
            logger.error('Named pipe connection error: %s (%s)' % (str(e), self.__class__))


class RemoteStdOutPipe(Pipes):
    def __init__(self, transport, pipe, permisssions):
        Pipes.__init__(self, transport, pipe, permisssions)

    def run(self):
        self.connectPipe()

        while True:
            try:
                ans = self.server.readFile(self.tid, self.fid, 0, 1024)
            except Exception as e:
                pass
            else:
                try:
                    global LastDataSent

                    if ans != LastDataSent:
                        sys.stdout.write(ans.decode('cp437'))
                        sys.stdout.flush()
                    else:
                        # Don't echo what I sent, and clear it up
                        LastDataSent = ''

                    # Just in case this got out of sync, i'm cleaning it up if there are more than 10 chars,
                    # it will give false positives tho.. we should find a better way to handle this.
                    if LastDataSent > 10:
                        LastDataSent = ''
                except Exception as e:
                    pass


class RemoteStdErrPipe(Pipes):
    def __init__(self, transport, pipe, permisssions):
        Pipes.__init__(self, transport, pipe, permisssions)

    def run(self):
        self.connectPipe()

        while True:
            try:
                ans = self.server.readFile(self.tid, self.fid, 0, 1024)
            except Exception as e:
                pass
            else:
                try:
                    sys.stderr.write(ans.decode('cp437'))
                    sys.stderr.flush()
                except Exception as e:
                    pass


class RemoteShell(cmd.Cmd):
    def __init__(self, server, port, credentials, tid, fid, share):
        cmd.Cmd.__init__(self, False)
        self.prompt = '\x08'
        self.server = server
        self.tid = tid
        self.fid = fid
        self.credentials = credentials
        self.share = share
        self.port = port

    def cmdloop(self):
        try:
            cmd.Cmd.cmdloop(self)
        except SessionError as e:
            # traceback.print_exc()
            logger.error('SMB error: %s' % (e.getErrorString(),))
        except NetBIOSTimeout as e:
            logger.error('SMB connection timed out')
        except keimpxError as e:
            logger.error(e)
        except KeyboardInterrupt as _:
            print()
            logger.info('User aborted')
            exit()
        except Exception as e:
            # traceback.print_exc()
            logger.error(str(e))

    def emptyline(self):
        self.send_data('\r\n')
        return

    def do_shell(self, command):
        process = Popen(command, shell=True, stdout=PIPE, stderr=STDOUT)
        stdout, _ = process.communicate()

        if stdout is not None:
            print(stdout)

        self.send_data('\r\n')

    def do_exit(self, line):
        self.send_data('exit\r\n')
        return

    def default(self, line=''):
        if isinstance(line, bytearray):
            self.send_data(line.decode('cp437') + '\r\n')
        else:
            self.send_data(line + '\r\n')

    def send_data(self, data, hideOutput=True):
        if hideOutput is True:
            global LastDataSent
            LastDataSent = data
        else:
            LastDataSent = ''

        self.server.writeFile(self.tid, self.fid, data)


class RemoteStdInPipe(Pipes):
    def __init__(self, transport, pipe, permisssions, share=None):
        Pipes.__init__(self, transport, pipe, permisssions, share)

    def run(self):
        self.connectPipe()
        self.remote_shell = RemoteShell(self.server, self.port, self.credentials, self.tid, self.fid, self.share)
        self.remote_shell.cmdloop()
