from bootstrapvz.base import Task
from .. import phases
from ..tools import log_check_call
import os.path
from . import assets
import initd


class AddOpenSSHPackage(Task):
	description = 'Adding openssh package'
	phase = phases.preparation

	@classmethod
	def run(cls, info):
		info.packages.add('openssh-server')


class AddSSHKeyGeneration(Task):
	description = 'Adding SSH private key generation init scripts'
	phase = phases.system_modification
	successors = [initd.InstallInitScripts]

	@classmethod
	def run(cls, info):
		init_scripts_dir = os.path.join(assets, 'init.d')
		install = info.initd['install']
		from subprocess import CalledProcessError
		try:
			log_check_call(['chroot', info.root,
			                'dpkg-query', '-W', 'openssh-server'])
			from bootstrapvz.common.releases import squeeze
			if info.manifest.release == squeeze:
				install['generate-ssh-hostkeys'] = os.path.join(init_scripts_dir, 'squeeze/generate-ssh-hostkeys')
			else:
				install['generate-ssh-hostkeys'] = os.path.join(init_scripts_dir, 'generate-ssh-hostkeys')
		except CalledProcessError:
			import logging
			logging.getLogger(__name__).warn('The OpenSSH server has not been installed, '
			                                 'not installing SSH host key generation script.')


class DisableSSHPasswordAuthentication(Task):
	description = 'Disabling SSH password authentication'
	phase = phases.system_modification

	@classmethod
	def run(cls, info):
		from ..tools import sed_i
		sshd_config_path = os.path.join(info.root, 'etc/ssh/sshd_config')
		sed_i(sshd_config_path, '^#PasswordAuthentication yes', 'PasswordAuthentication no')


class EnableRootLogin(Task):
	description = 'Enabling SSH login for root'
	phase = phases.system_modification

	@classmethod
	def run(cls, info):
		sshdconfig_path = os.path.join(info.root, 'etc/ssh/sshd_config')
		if os.path.exists(sshdconfig_path):
			from bootstrapvz.common.tools import sed_i
			sed_i(sshdconfig_path, '^PermitRootLogin .*', 'PermitRootLogin yes')
		else:
			import logging
			logging.getLogger(__name__).warn('The OpenSSH server has not been installed, '
			                                 'not enabling SSH root login.')


class DisableRootLogin(Task):
	description = 'Disabling SSH login for root'
	phase = phases.system_modification

	@classmethod
	def run(cls, info):
		sshdconfig_path = os.path.join(info.root, 'etc/ssh/sshd_config')
		if os.path.exists(sshdconfig_path):
			from bootstrapvz.common.tools import sed_i
			sed_i(sshdconfig_path, '^PermitRootLogin .*', 'PermitRootLogin no')
		else:
			import logging
			logging.getLogger(__name__).warn('The OpenSSH server has not been installed, '
			                                 'not disabling SSH root login.')


class DisableSSHDNSLookup(Task):
	description = 'Disabling sshd remote host name lookup'
	phase = phases.system_modification

	@classmethod
	def run(cls, info):
		sshd_config_path = os.path.join(info.root, 'etc/ssh/sshd_config')
		with open(sshd_config_path, 'a') as sshd_config:
			sshd_config.write('UseDNS no')


class ShredHostkeys(Task):
	description = 'Securely deleting ssh hostkeys'
	phase = phases.system_cleaning

	@classmethod
	def run(cls, info):
		ssh_hostkeys = ['ssh_host_dsa_key',
		                'ssh_host_rsa_key']
		from bootstrapvz.common.releases import wheezy
		if info.manifest.release >= wheezy:
			ssh_hostkeys.append('ssh_host_ecdsa_key')

		private = [os.path.join(info.root, 'etc/ssh', name) for name in ssh_hostkeys]
		public = [path + '.pub' for path in private]

		from ..tools import log_check_call
		log_check_call(['shred', '--remove'] + private + public)
