#!/usr/bin/python2.7
# -*- coding: utf-8 -*-
"""
Run various powershell commands on windows clients via winrm
"""
#
# Copyright (C) 2010-2017 Univention GmbH
#
# http://www.univention.de/
#
# All rights reserved.
#
# The source code of this program is made available
# under the terms of the GNU Affero General Public License version 3
# (GNU AGPL V3) as published by the Free Software Foundation.
#
# Binary versions of this program provided by Univention to you as
# well as other copyrighted, protected or trademarked materials like
# Logos, graphics, fonts, specific documentations and configurations,
# cryptographic keys etc. are subject to a license agreement between
# you and Univention and not subject to the GNU AGPL V3.
#
# In the case you use this program under the terms of the GNU AGPL V3,
# the program is provided in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License with the Debian GNU/Linux or Univention distribution in file
# /usr/share/common-licenses/AGPL-3; if not, see
# <http://www.gnu.org/licenses/>.

import sys
import os
import argparse
import winrm
import socket
import time
import ConfigParser
import importlib
import glob
import logging
import functools

COMMAND_DIR = None
for path in [
	os.path.join(os.environ['HOME'], 'git/ucs-ec2-tools/shared-utils/ucs-winrm-plugins'),
	'./tools/shared-utils/ucs-winrm-plugins',
	'./shared-utils/ucs-winrm-plugins',
	'./ucs-winrm-plugins',
	'/root/shared-utils/ucs-winrm-plugins']:
	if os.path.exists(path):
		COMMAND_DIR = path
		break
assert COMMAND_DIR
sys.path.append(COMMAND_DIR)

CONFIG_FILE = os.path.join(os.environ['HOME'], '.ucs-winrm.ini')


class UcsWinFail(Exception):
	pass


class RedirectStdOut(object):

	def __init__(self, stdout=None):
		self._stdout = stdout or sys.stdout

	def __enter__(self):
		self.old_stdout = sys.stdout
		self.old_stdout.flush()
		sys.stdout = self._stdout

	def __exit__(self, exc_type, exc_value, traceback):
		self._stdout.flush()
		sys.stdout = self.old_stdout


def check_dryrun(f):
	def wrapped(self, *args, **kwargs):
		self.logger.debug('-> {name}()'.format(name=f.func_name))
		for arg in args:
			for i in arg.split('\n'):
				self.logger.debug('    {i}'.format(i=i))
		if self.args.dry_run:
			return 0
		return f(self, *args, **kwargs)
	return wrapped


class Plugins(object):
	pass


class UcsWin(object):

	PS = '''
		Set-ExecutionPolicy Unrestricted
		$ErrorActionPreference = "Stop"
		$content = [IO.File]::ReadAllText("%(psfile)s");
		$code = [System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($content))
		[System.Threading.Thread]::CurrentThread.CurrentCulture = 'en-US';
		[System.Threading.Thread]::CurrentThread.CurrentUICulture = 'en-US';
		Invoke-Expression $code
	'''

	PS_CREDSSP = '''
		#Enable-WSManCredSSP -Role Server -Force
		#Enable-WSManCredSSP -Role Client -DelegateComputer $env:computername -Force
		Set-ExecutionPolicy Unrestricted
		$strUser = '%(domain)s\%(run_as_user)s'
		$strPass = ConvertTo-SecureString -String "%(run_as_password)s" -AsPlainText -Force
		$objCred = New-Object -TypeName System.Management.Automation.PSCredential -ArgumentList ($strUser, $strPass)
		$content = [IO.File]::ReadAllText("%(psfile)s");
		[System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($content)) > out.ps1
		Invoke-Command -Authentication Credssp -ComputerName $env:computername -Credential $objCred -FilePath "out.ps1"
	'''

	# this approach uses the LogonUser user dll function to Impersonate the current context
	PS_IMPERSONATE = '''

		Set-ExecutionPolicy Unrestricted

		$ErrorActionPreference = "Stop"
		$content = [IO.File]::ReadAllText("%(psfile)s");
		$code = [System.Text.Encoding]::UTF8.GetString([System.Convert]::FromBase64String($content))
		[System.Threading.Thread]::CurrentThread.CurrentCulture = 'en-US';
		[System.Threading.Thread]::CurrentThread.CurrentUICulture = 'en-US';

# Import the LogonUser Function from advapi32.dll and the CloseHandle Function from kernel32.dll
Add-Type -Namespace Import -Name Win32 -MemberDefinition @'
	[DllImport("advapi32.dll", SetLastError = true)]
	public static extern bool LogonUser(string user, string domain, string password, int logonType, int logonProvider, out IntPtr token);
	[DllImport("kernel32.dll", SetLastError = true)]
	public static extern bool CloseHandle(IntPtr handle);
'@

		$Username = "%(run_as_user)s"
		$Domain = "%(domain)s"
		$Password = "%(run_as_password)s"
		$Global:ImpersonatedUser = @{}
		$tokenHandle = 0
		$returnValue = [Import.Win32]::LogonUser($Username, $Domain, $Password, 2, 0, [ref]$tokenHandle)
		if (!$returnValue) {
			$errCode = [System.Runtime.InteropServices.Marshal]::GetLastWin32Error();
			Write-Host "Impersonate-User failed a call to LogonUser with error code: $errCode"
			throw [System.ComponentModel.Win32Exception]$errCode
		}

		# Call the Impersonate method with the returned token. An ImpersonationContext is returned and stored in the
		# Global variable so that it may be used after script run.
		$Global:ImpersonatedUser.ImpersonationContext = [System.Security.Principal.WindowsIdentity]::Impersonate($tokenHandle)

		# Close the handle to the token. Voided to mask the Boolean return value.
		[void][Import.Win32]::CloseHandle($tokenHandle)

		Invoke-Expression $code
		'''

	def __init__(self):

		# read config file
		self.config = ConfigParser.ConfigParser()
		self.config.read(CONFIG_FILE)

		# parser
		self.parser = argparse.ArgumentParser(description=sys.modules[__name__].__doc__)
		self.parser.usage = '%(prog)s COMMAND [OPTIONS]'
		self.subparsers = self.parser.add_subparsers(description='type %(prog)s <action> --help for further help and possible arguments', metavar='commands')

		# load plugins
		self.plugins = Plugins()
		for name in sorted(glob.glob(COMMAND_DIR + '/[!_]*.py')):
			name, ext = os.path.splitext(os.path.basename(name))
			module = importlib.import_module(name)
			if hasattr(module, 'name') and hasattr(module, 'description'):
				# global options
				subparser = self.subparsers.add_parser(module.name, description=module.description, help=module.description)
				subparser.set_defaults(module=module)
				subparser.add_argument('--client', default=self._get_default('client'), help='The IP of the windows client')
				subparser.add_argument('--user', default=self._get_default('user'), help='username')
				subparser.add_argument('--password', default=self._get_default('password'), help='password')
				subparser.add_argument('--debug', action='store_true', default=self._get_default('debug', default=False), help='enable debug')
				subparser.add_argument('--dry-run', action='store_true', default=self._get_default('debug', default=False), help='dry run')
				subparser.add_argument("--loglevel", type=str, default='debug', help="log level", choices=['debug', 'error', 'info', 'warn'])
				subparser.add_argument("--credssp", action='store_true', help="use Invoke-Command -Authentication Credssp on the client to run command")
				subparser.add_argument("--impersonate", action='store_true', help="use Impersonate method to run command")
				subparser.add_argument('--domain', default=self._get_default('domain'), help='the windows domain name')
				subparser.add_argument('--run-as-user', default=self._get_default('run_as_user'), help='the user name for the credssp/impersonate')
				subparser.add_argument('--run-as-password', default=self._get_default('run_as_password'), help='the password for the credssp/impersonate')
				# command options
				if hasattr(module, 'args'):
					for arg in module.args:
						if 'default' not in module.args[arg]:
							module.args[arg]['default'] = self._get_default(arg)
						subparser.add_argument('--%s' % arg, **module.args[arg])
				setattr(self.plugins, name, functools.partial(self._run_module, module))

		self.args = self.parser.parse_args()
		if self.args.domain:
			self.args.nb_domain = self.args.domain.split('.')[0]
		if self.args.run_as_password is None:
			self.args.run_as_password = self.args.password
		if self.args.run_as_user is None:
			self.args.run_as_user = self.args.user
		self.last_stdout = None

		# logging
		self.logger = logging.getLogger('ucs-winrm')
		formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
		self.logger = logging.getLogger('ucs-winrm')
		handler = logging.StreamHandler()
		handler.setFormatter(formatter)
		self.logger.addHandler(handler)
		if self.args.loglevel == 'debug':
			self.logger.setLevel(logging.DEBUG)
		elif self.args.loglevel == 'warn':
			self.logger.setLevel(logging.WARN)
		elif self.args.loglevel == 'error':
			self.logger.setLevel(logging.ERROR)
		elif self.args.loglevel == 'info':
			self.logger.setLevel(logging.INFO)
		else:
			self.logger.setLevel(logging.DEBUG)

		self.logger.debug('-> {action}'.format(action=self.args.module.name))
		for arg in vars(self.args):
			self.logger.debug('   arg {arg}={value}'.format(arg=arg, value=vars(self.args)[arg]))
		self._run_module(self.args.module)
		sys.exit(0)

	def _run_module(self, module, **kwargs):
		# execute pre hook
		if hasattr(module, 'pre'):
			module.pre(self)
		if hasattr(module, 'ps'):
			cmd = module.ps % kwargs if kwargs else module.ps % vars(self.args)
			self.run(cmd)
		# execute post hook
		if hasattr(module, 'post'):
			module.post(self)

	def _get_default(self, parameter, default=None):
		val = None
		# check env variable
		env = 'WINRM_' + parameter.replace('-', '_').upper()
		val = os.environ.get(env)
		if self.config.has_section('default') and self.config.has_option('default', parameter):
			val = self.config.get('default', parameter)
		if val is None:
			val = default
		return val

	##########
	# helper #
	##########

	def copy_file(self, session, remote_file, data=None, filename=None):
		# TODO hanlde data
		if filename:
			data = filename
		chunks = 2500
		b64 = data.encode('base64').replace('\n', '')
		# remove remote file
		with RedirectStdOut(stdout=open(os.devnull, 'w')):
			session.run_ps('If (Test-Path "%(rfile)s") { Remove-Item "%(rfile)s" }' % dict(rfile=remote_file))
		for i in range(0, len(b64), chunks):
			r = session.run_ps('[System.IO.File]::AppendAllText("%s", "%s")' % (remote_file, b64[i:i + chunks]))
			if r.status_code != 0:
				raise UcsWinFail('failed to copy command \ncode: %s\nout: %s\nerr: %s' % (r.status_code, r.std_out, r.std_err))

	@check_dryrun
	def run(self, command, silent=False):
		s = winrm.Session(self.args.client, auth=(self.args.user, self.args.password), transport='plaintext')
		# run_ps hits the windows Maximum Length of Command Line String
		psfile = 'c:\\.tmp.winrm-ucs.ps1'
		self.copy_file(s, psfile, data=command)
		if self.args.credssp and vars(self.args).get('domain'):
			self.args.psfile = psfile
			cmd = self.PS_CREDSSP % vars(self.args)
		elif self.args.impersonate and self.args.domain:
			self.args.psfile = psfile
			cmd = self.PS_IMPERSONATE % vars(self.args)
		else:
			cmd = self.PS % dict(psfile=psfile)
		with RedirectStdOut(stdout=open(os.devnull, 'w')):
			r = s.run_ps(cmd)
		if r.status_code != 0:
			raise UcsWinFail('command %s failed with\ncode: %s\nout: %s\nerr: %s' % (command, r.status_code, r.std_out, r.std_err))
		if r.std_out:
			if not silent:
				print(r.std_out)
			self.last_stdout = r.std_out

	@check_dryrun
	def client_reachable(self):
		try:
			self.logger.debug('check if %s is available' % self.args.client)
			s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
			s.settimeout(2)
			s.connect((self.args.client, 5985))
			s = winrm.Session(self.args.client, auth=(self.args.user, self.args.password), transport='plaintext')
			s.run_cmd('dir')
			return True
		except socket.error:
			return False
		except winrm.exceptions.WinRMTransportError:
			return False

	@check_dryrun
	def wait_until_client_is_gone(self, timeout=120):
		for i in range(timeout):
			if not self.client_reachable():
				self.logger.debug('%s is gone' % self.args.client)
				return True
			time.sleep(1)
		raise UcsWinFail('wait_until_client_is_gone for %s failed (timeout=%s)' % (self.args.client, timeout))

	@check_dryrun
	def wait_for_client(self, timeout=300):
		for i in range(timeout):
			if self.client_reachable():
				self.logger.debug('%s is available' % self.args.client)
				return True
			time.sleep(1)
		raise UcsWinFail('wait_for_client for %s failed (timeout=%s)' % (self.args.client, timeout))

	def error(self, message):
		raise UcsWinFail(message)


if __name__ == '__main__':
	UcsWin()
