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

from __future__ import print_function

import ntpath
import os
import shlex
import sys

from lib.avservices import AVSERVICES
from lib.common import DataStore, is_local_admin, SMBServer
from lib.exceptions import missingPermission
from lib.logger import logger
from lib.smbexec import SvcShell

try:
    from impacket.dcerpc.v5 import scmr
    from impacket.dcerpc.v5.dtypes import NULL
    from impacket.smbconnection import SessionError
    from impacket.crypto import encryptSecret
except ImportError:
    sys.stderr.write('services: 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 services.py example #
#################################################################
class SvcCtl(object):
    def __init__(self):
        pass

    def services(self, srvname):
        self.__scmr_connect()
        self.__scmr_list(srvname)
        self.__scmr_disconnect()

    def status(self, srvname, return_state=False):
        self.__scmr_connect()
        self.__scmr_srv_manager(srvname)
        resp = self.__scmr_state(srvname, return_state)
        self.__scmr_disconnect()

        return resp

    def query(self, srvname, return_answer=False):
        self.__scmr_connect()
        self.__scmr_srv_manager(srvname)
        resp = self.__scmr_config(srvname, return_answer)
        self.__scmr_disconnect()

        return resp

    def start(self, srvname, srvargs=''):
        self.__scmr_connect()
        self.__scmr_srv_manager(srvname)
        self.__scmr_start(srvname, srvargs)
        self.__scmr_disconnect(srvname)

    def stop(self, srvname):
        self.__scmr_connect()
        self.__scmr_srv_manager(srvname)
        self.__scmr_stop(srvname)
        self.__scmr_disconnect(srvname)

    def change(self, srvname, display=None, path=None, service_type=None, start_type=None, start_name=None,
               password=None):
        self.__scmr_connect()
        self.__scmr_srv_manager(srvname)
        self.__scmr_change(display, path, service_type, start_type, start_name, password)
        self.__scmr_disconnect(srvname)

    def deploy(self, srvname, local_file=None, srvargs='', remote_file=None, displayname=None):
        self.oldpwd = self.pwd
        self.pwd = '\\'

        self.__scmr_bin_upload(local_file, remote_file)
        self.__scmr_connect()
        self.__scmr_create(srvname, remote_file, displayname)
        self.__scmr_srv_manager(srvname)
        self.__scmr_start(srvname, srvargs)
        self.__scmr_disconnect(srvname)

        self.pwd = self.oldpwd

    def undeploy(self, srvname):
        self.oldpwd = self.pwd
        self.pwd = '\\'

        self.__scmr_connect()
        self.__scmr_srv_manager(srvname)
        resp = scmr.hRQueryServiceConfigW(self.__rpc, self.__service_handle)
        remote_file = resp['lpServiceConfig']['lpBinaryPathName'][:-1]
        remote_file = str(os.path.basename(remote_file.replace('\\', '/')))

        if self.__scmr_state(srvname, return_state=True) == 'RUNNING':
            self.__scmr_stop(srvname)

        self.__scmr_delete(srvname)
        self.__scmr_disconnect(srvname)
        self.__scmr_bin_remove(remote_file)
        self.pwd = self.oldpwd

    def stopav(self):
        self.__scmr_connect()
        resp = scmr.hREnumServicesStatusW(self.__rpc, self.__mgr_handle, dwServiceState=scmr.SERVICE_STATE_ALL)

        for i in range(len(resp)):
            name = resp[i]['lpServiceName'][:-1]
            display = resp[i]['lpDisplayName'][:-1]
            state = resp[i]['ServiceStatus']['dwCurrentState']
            if state == scmr.SERVICE_RUNNING and name in AVSERVICES:
                try:
                    self.stop(name)
                except:
                    print("Couldn't stop %s (%s), currently in state: %s" % (
                        display, name, self.__scmr_parse_state(state)))

    def svcexec(self, command, mode='SHARE', display=True):
        if mode == 'SERVER' and not is_local_admin():
            err = ("keimpx needs to be run as Administrator/root to use svcshell. "
                   "Privileged port is needed to run SMB server.")
            raise missingPermission(err)

        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])

        self.__scmr_connect()

        try:
            if mode == 'SERVER':
                self.__serverThread = SMBServer(self.smbserver_share)
                self.__serverThread.daemon = True
                self.__serverThread.start()

            if os.path.exists(command_and_args[0]):
                command = ntpath.join(DataStore.share_path, os.path.basename(command))

            self.svc_shell = SvcShell(self.__rpc, self.__mgr_handle, self.trans, self.smbserver_share, mode, display)
            self.svc_shell.onecmd(command)

            if mode == 'SERVER':
                self.__serverThread.stop()
        except SessionError as e:
            # traceback.print_exc()
            logger.error('SMB error: %s' % (e.getErrorString(),))
        except KeyboardInterrupt as _:
            print()
            logger.info('User aborted')
        except Exception as e:
            # traceback.print_exc()
            logger.error(str(e))

        sys.stdout.flush()
        self.__scmr_disconnect()

        if os.path.exists(command_and_args[0]):
            self.rm(os.path.basename(command_and_args[0]))

    def svcshell(self, mode='SHARE'):
        if mode == 'SERVER' and not is_local_admin():
            err = ("keimpx needs to be run as Administrator/root to use svcshell. Privileged port is needed to run SMB "
                   "server.")
            raise missingPermission(err)

        self.__scmr_connect()

        try:
            if mode == 'SERVER':
                self.__serverThread = SMBServer(self.smbserver_share)
                self.__serverThread.daemon = True
                self.__serverThread.start()

            self.svc_shell = SvcShell(self.__rpc, self.__mgr_handle, self.trans, self.smbserver_share, mode)
            self.svc_shell.cmdloop()

            if mode == 'SERVER':
                self.__serverThread.stop()
        except SessionError as e:
            # traceback.print_exc()
            logger.error('SMB error: %s' % (e.getErrorString(),))
        except KeyboardInterrupt as _:
            print()
            logger.info('User aborted')
        except Exception as e:
            # import traceback
            # traceback.print_exc()
            logger.error(str(e))

        sys.stdout.flush()
        self.__scmr_disconnect()

    def __scmr_srv_manager(self, srvname):
        self.__resp = scmr.hROpenServiceW(self.__rpc, self.__mgr_handle, '%s\x00' % srvname)
        self.__service_handle = self.__resp['lpServiceHandle']

    def __scmr_connect(self):
        '''
        Connect to svcctl named pipe
        '''
        self.smb_transport('svcctl')

        self.__dce = self.trans.get_dce_rpc()
        self.__dce.bind(scmr.MSRPC_UUID_SCMR)
        self.__rpc = self.__dce
        self.__resp = scmr.hROpenSCManagerW(self.__dce)
        self.__mgr_handle = self.__resp['lpScHandle']

    def __scmr_disconnect(self, srvname=None):
        '''
        Disconnect from svcctl named pipe
        '''
        if srvname:
            scmr.hRCloseServiceHandle(self.__rpc, self.__service_handle)

        if self.__mgr_handle:
            scmr.hRCloseServiceHandle(self.__rpc, self.__mgr_handle)

        self.__dce.disconnect()

    def __scmr_bin_upload(self, local_file, remote_file):
        '''
        Upload the service executable
        '''
        self.use(DataStore.writable_share)
        self.__pathname = ntpath.join(DataStore.writable_share, remote_file)
        logger.info('Uploading the service executable to %s' % self.__pathname)
        self.upload(local_file, remote_file)

    def __scmr_bin_remove(self, remote_file):
        '''
        Remove the service executable
        '''
        self.use(DataStore.writable_share)
        self.__pathname = ntpath.join(DataStore.writable_share, remote_file)
        logger.info('Removing the service executable %s' % self.__pathname)
        self.rm(remote_file)

    def __scmr_create(self, srvname, remote_file, displayname=None):
        '''
        Create the service
        '''
        logger.info('Creating the service %s' % srvname)

        if not displayname:
            displayname = srvname

        self.__pathname = ntpath.join(DataStore.share_path, remote_file)
        scmr.hRCreateServiceW(self.__rpc, self.__mgr_handle, '%s\x00' % srvname, '%s\x00' % displayname,
                              lpBinaryPathName='%s\x00' % self.__pathname, dwStartType=scmr.SERVICE_DEMAND_START)

    def __scmr_delete(self, srvname):
        '''
        Delete the service
        '''
        logger.info('Deleting the service %s' % srvname)
        scmr.hRDeleteService(self.__rpc, self.__service_handle)

    def __scmr_parse_config(self, resp):
        print('TYPE              : %2d - ' % resp['lpServiceConfig']['dwServiceType'])

        if resp['lpServiceConfig']['dwServiceType'] & 0x1:
            print('SERVICE_KERNEL_DRIVER')
        if resp['lpServiceConfig']['dwServiceType'] & 0x2:
            print('SERVICE_FILE_SYSTEM_DRIVER')
        if resp['lpServiceConfig']['dwServiceType'] & 0x10:
            print('SERVICE_WIN32_OWN_PROCESS')
        if resp['lpServiceConfig']['dwServiceType'] & 0x20:
            print('SERVICE_WIN32_SHARE_PROCESS')
        if resp['lpServiceConfig']['dwServiceType'] & 0x100:
            print('SERVICE_INTERACTIVE_PROCESS')

        print('START_TYPE        : %2d - ' % resp['lpServiceConfig']['dwStartType'])

        if resp['lpServiceConfig']['dwStartType'] == 0x0:
            print('BOOT START')
        elif resp['lpServiceConfig']['dwStartType'] == 0x1:
            print('SYSTEM START')
        elif resp['lpServiceConfig']['dwStartType'] == 0x2:
            print('AUTO START')
        elif resp['lpServiceConfig']['dwStartType'] == 0x3:
            print('DEMAND START')
        elif resp['lpServiceConfig']['dwStartType'] == 0x4:
            print('DISABLED')
        else:
            print('UNKNOWN')

        print('ERROR_CONTROL     : %2d - ' % resp['lpServiceConfig']['dwErrorControl'])

        if resp['lpServiceConfig']['dwErrorControl'] == 0x0:
            print('IGNORE')
        elif resp['lpServiceConfig']['dwErrorControl'] == 0x1:
            print('NORMAL')
        elif resp['lpServiceConfig']['dwErrorControl'] == 0x2:
            print('SEVERE')
        elif resp['lpServiceConfig']['dwErrorControl'] == 0x3:
            print('CRITICAL')
        else:
            print('UNKNOWN')

        print('BINARY_PATH_NAME  : %s' % resp['lpServiceConfig']['lpBinaryPathName'][:-1])
        print('LOAD_ORDER_GROUP  : %s' % resp['lpServiceConfig']['lpLoadOrderGroup'][:-1])
        print('TAG               : %d' % resp['lpServiceConfig']['dwTagId'])
        print('DISPLAY_NAME      : %s' % resp['lpServiceConfig']['lpDisplayName'][:-1])
        print('DEPENDENCIES      : %s' % resp['lpServiceConfig']['lpDependencies'][:-1])
        print('SERVICE_START_NAME: %s' % resp['lpServiceConfig']['lpServiceStartName'][:-1])

    def __scmr_parse_state(self, state):
        if state == scmr.SERVICE_CONTINUE_PENDING:
            return 'CONTINUE PENDING'
        elif state == scmr.SERVICE_PAUSE_PENDING:
            return 'PAUSE PENDING'
        elif state == scmr.SERVICE_PAUSED:
            return 'PAUSED'
        elif state == scmr.SERVICE_RUNNING:
            return 'RUNNING'
        elif state == scmr.SERVICE_START_PENDING:
            return 'START PENDING'
        elif state == scmr.SERVICE_STOP_PENDING:
            return 'STOP PENDING'
        elif state == scmr.SERVICE_STOPPED:
            return 'STOPPED'
        else:
            return 'UNKOWN'

    def __scmr_state(self, srvname, return_state=False):
        '''
        Display state of a service
        '''
        logger.info('Querying the state of service %s' % srvname)

        resp = scmr.hRQueryServiceStatus(self.__rpc, self.__service_handle)
        state = resp['lpServiceStatus']['dwCurrentState']

        if return_state:
            return self.__scmr_parse_state(state)
        else:
            print('Service %s state is: %s' % (srvname, self.__scmr_parse_state(state)))

    def __scmr_config(self, srvname, return_answer=False):
        '''
        Display a service configuration
        '''
        logger.info('Querying the service configuration of service %s' % srvname)

        resp = scmr.hRQueryServiceConfigW(self.__rpc, self.__service_handle)

        if return_answer:
            return resp

        print('Service %s information:' % srvname)
        self.__scmr_parse_config(resp)

    def __scmr_start(self, srvname, srvargs=''):
        '''
        Start the service
        '''
        logger.info('Starting the service %s' % srvname)

        if srvargs:
            srvargs = str(srvargs).split(' ')
        else:
            srvargs = []

        scmr.hRStartServiceW(self.__rpc, self.__service_handle, argc=len(srvargs), argv=srvargs)
        self.__scmr_state(srvname)

    def __scmr_stop(self, srvname):
        '''
        Stop the service
        '''
        logger.info('Stopping the service %s' % srvname)

        scmr.hRControlService(self.__rpc, self.__service_handle, scmr.SERVICE_CONTROL_STOP)
        self.__scmr_state(srvname)

    def __scmr_change(self, display=None, path=None, service_type=None, start_type=None, start_name=None,
                      password=None):
        '''
        Change the configuration of a service
        '''
        if start_type is not None:
            start_type = int(start_type)
        else:
            start_type = scmr.SERVICE_NO_CHANGE

        if service_type is not None:
            service_type = int(service_type)
        else:
            service_type = scmr.SERVICE_NO_CHANGE

        if display is not None:
            display = '%s\x00' % display
        else:
            display = NULL

        if path is not None:
            path = '%s\x00' % path
        else:
            path = NULL

        if start_name is not None:
            start_name = '%s\x00' % start_name
        else:
            start_name = NULL

        if password is not None:
            s = self.trans.get_smb_connection()
            key = s.getSessionKey()
            password = ('%s\x00' % password).encode('utf-16le')
            password = encryptSecret(key, password)
        else:
            password = NULL

        scmr.hRChangeServiceConfigW(self.__rpc, self.__service_handle, service_type, start_type,
                                    scmr.SERVICE_ERROR_IGNORE, path, NULL, NULL, NULL, 0, start_name, password, 0,
                                    display)

    def __scmr_list_parse(self, srvname, resp):
        '''
        Parse list of services
        '''
        services = []

        for i in range(len(resp)):
            name = resp[i]['lpServiceName'][:-1]
            display = resp[i]['lpDisplayName'][:-1]
            state = resp[i]['ServiceStatus']['dwCurrentState']

            if srvname:
                srvname = srvname.strip('*')

                if srvname.lower() not in display.lower() and srvname.lower() not in name.lower():
                    continue

            services.append((display, name, state))

        services.sort()

        for service in services:
            print('%s (%s): %-80s' % (service[0], service[1], self.__scmr_parse_state(service[2])))

        return len(services)

    def __scmr_list(self, srvname):
        '''
        List services
        '''
        logger.info('Listing services')

        resp = scmr.hREnumServicesStatusW(self.__rpc, self.__mgr_handle, dwServiceState=scmr.SERVICE_STATE_ALL)
        num = self.__scmr_list_parse(srvname, resp)

        print('\nTotal services: %d\n' % num)
