from contextlib import contextmanager

from parallels.common.actions.hosting_settings.convert.ppa.dns \
	import DNS as ActionHostingSettingsConvertPPADNS
from parallels.expand_migrator.data_source import DataSource

class DNS(ActionHostingSettingsConvertPPADNS):
	def _get_all_source_ips(self, global_context, subscription):
		"""Get all IPs of all source servers."""
		with self._get_data_source(global_context) as ds:
			return (
				{global_context.conn.expand.conn_settings.ip} | \
				ds.get_dns_server_ips() | \
				ds.get_plesk_server_ips() | \
				ds.get_mail_server_ips() | \
				ds.get_db_server_ips()
			)

	@contextmanager
	def _get_data_source(self, global_context):
		with global_context.conn.get_expand_runner() as expand_runner:
			yield DataSource(global_context.conn.expand.api(), expand_runner)
