resolve-hostnames

#!/usr/bin/env python3

import os, sys, re, socket, logging


class LogMessage:
	def __init__(self, fmt, a, k): self.fmt, self.a, self.k = fmt, a, k
	def __str__(self): return self.fmt.format(*self.a, **self.k) if self.a or self.k else self.fmt

class LogStyleAdapter(logging.LoggerAdapter):
	def __init__(self, logger, extra=None):
		super(LogStyleAdapter, self).__init__(logger, extra or {})
	def log(self, level, msg, *args, **kws):
		if not self.isEnabledFor(level): return
		log_kws = {} if 'exc_info' not in kws else dict(exc_info=kws.pop('exc_info'))
		msg, kws = self.process(msg, kws)
		self.logger._log(level, LogMessage(msg, args, kws), (), **log_kws)

get_logger = lambda name: LogStyleAdapter(logging.getLogger(name))


class AddressError(Exception): pass


def sock_enum_val(name_or_num, prefix=None):
	name_or_num = name_or_num.upper()
	if prefix: prefix = prefix.upper()
	try:
		return getattr( socket, name_or_num
			if not prefix else '{}_{}'.format(prefix, name_or_num) )
	except AttributeError: pass
	try: return int(name_or_num)
	except ValueError: pass
	raise AddressError(( 'Failed to resolve'
		' socket parameter name/value: {!r}' ).format(name_or_num))


def get_socket_info( host, port=0, family=0,
		socktype=0, protocol=0, force_unique_address=False ):
	log_params = [port, family, socktype, protocol]
	log.debug('Resolving: {} (params: {})', host, log_params)
	try:
		addrinfo = socket.getaddrinfo(host, port, family, socktype, protocol)
		if not addrinfo: raise socket.gaierror('No addrinfo for host: {}'.format(host))
	except (socket.gaierror, socket.error) as err:
		raise AddressError( 'Failed to resolve host:'
			' {!r} (params: {}) - {} {}'.format(host, log_params, type(err), err) )

	ai_af, ai_addr = set(), list()
	for family, _, _, hostname, addr in addrinfo:
		ai_af.add(family)
		ai_addr.append((addr[0], family))

	if len(ai_af) > 1:
		af_names = dict((v, k) for k,v in vars(socket).items() if k.startswith('AF_'))
		ai_af_names = list(af_names.get(af, str(af)) for af in ai_af)
		if socket.AF_INET not in ai_af:
			log.fatal(
				'Ambiguous socket host specification (matches address famlies: {}),'
					' refusing to pick one at random - specify socket family instead. Addresses: {}'
				', '.join(ai_af_names), ', '.join(ai_addr) )
			raise AddressError
		log.warn( 'Specified host matches more than'
			' one address family ({}), using it as IPv4 (AF_INET).', ai_af_names )
		af = socket.AF_INET
	else: af = list(ai_af)[0]

	for addr, family in ai_addr:
		if family == af: break
	else: raise AddressError
	ai_addr_unique = set(ai_addr)
	if len(ai_addr_unique) > 1:
		if force_unique_address:
			raise AddressError('Address matches more than one host: {}'.format(ai_addr_unique))
		log.warn( 'Specified host matches more than'
			' one address ({}), using first one: {}', ai_addr_unique, addr )

	return addr, port


def main(args=None):
	import argparse
	parser = argparse.ArgumentParser(
		description='Resolve all provided hostnames (or hostname-patterns,'
				' may have port number in them) in the input text to their addresses.'
			' Result is dumped to stdout.' )

	parser.add_argument('hosts', nargs='*', metavar='host[:port]',
		help='Hosts to find/resolve in an input file/stream.'
			' Matched there only when separated by non-hostname char boundary,'
				' so that "host.tld" will be matched in "name host.tld;", but not in "mail.host.tld".')

	parser.add_argument('-f', '--file',
		metavar='path', help='File to resolve stuff in (default: use stdin).')
	parser.add_argument('-o', '--out',
		metavar='path', help='Output file or path (default: use stdout).')

	parser.add_argument('-p', '--port', type=int,
		metavar='number/service_name',
		help='Append the port to all passed hosts as "<host>:<port>".')
	parser.add_argument('-m', '--match', metavar='regexp',
		help='Regexp (python) to match host/port by, without ^ and $ matchers.'
			' Must contain either exactly one capture group (use "(?:...)" to avoid creating these),'
				' or named "name" group, and can also contain'
				' "port" and "protocol" named groups. Example: "(?P<name>.*\.(?:my\.tld1|my.tld2)):\d+":')
	parser.add_argument('-r', '--replace-group0', action='store_true',
		help='Replace not only "name" capture group, but whole regexp match (group 0).')

	parser.add_argument('-u', '--force-unique', action='store_true',
		help='Treat names that resolve to non-unique address as failures.'
			' This can be important for e.g. firewall configs, to avoid ambiguity.')
	parser.add_argument('-s', '--skip-unresolved',
		action='store_true', help='Do not stop on any unresolvable hosts.')

	parser.add_argument('-l', '--protocol', metavar='protocol_name/num', default='tcp',
		help='Name of a protocol to get address'
				' for (default: %(default)s), except if matched by --match.'
			' Should correspond directly to SOL_* names or be an integer. Empty or "0" = any.')
	parser.add_argument('-a', '--addr-family', metavar='family_name/num',
		help='Address family to resolve name to'
				' (inet = IPv4, inet6 = IPv6, etc; default: any).'
			' Should correspond directly to AF_* names or be an integer. Empty or "0" = any.')
	parser.add_argument('-t', '--sock-type', metavar='socket_type/num', default='stream',
		help='Socket type (stream, dgram, raw, etc; default: %(default)s).'
			' Should correspond directly to SOCK_* names or be an integer. Empty or "0" = any.')

	parser.add_argument('-d', '--debug', action='store_true', help='Verbose operation mode.')
	opts = parser.parse_args(sys.argv[1:] if args is None else args)

	if (opts.match and opts.hosts) or not (opts.match or opts.hosts):
		parser.error('Either --match or list of hosts should be specified (but not both).')

	global log
	logging.basicConfig(level=logging.DEBUG if opts.debug else logging.WARNING)
	log = get_logger('main')

	for k, pre in [('protocol', 'sol'), ('addr_family', 'af'), ('sock_type', 'sock')]:
		val = getattr(opts, k)
		val = 0 if not val else sock_enum_val(val, pre)
		setattr(opts, k, val)

	regexp = opts.match
	if not regexp:
		hc, names = r'[\w\d.-]', list()
		for name in opts.hosts:
			port = opts.port
			if re.search(r'^{}+:\d+$'.format(hc), name): name, port = name.rsplit(':', 1)
			else: assert re.search(r'^{}+$'.format(hc), name), name
			name = re.escape(name)
			if port: name += ':{}'.format(re.escape(bytes(port)))
			names.append(name)
		regexp = re.compile(
			r'(?<!{hc})(?P<name>{names})(?!{hc})'.format(hc=hc, names='|'.join(names)) )
	else: regexp = re.compile(regexp)

	src = open(opts.file).read() if opts.file else sys.stdin.read()
	src_diff = 0

	for match in regexp.finditer(src):
		# Get
		try:
			group, name = 1, match.groups()
			if len(name) != 1: group, name = 'name', match.group('name').rsplit(':', 1)
		except IndexError:
			log.error( 'Failed to get "name" group from'
				' regexp <re: {}> match: {!r}', opts.match, match.group(0) )
			return 1
		if opts.replace_group0: group = 0
		name, port = name if len(name) > 1 else (name[0], 0)
		try: protocol = match.group('protocol')
		except IndexError: protocol = opts.protocol

		# Resolve
		try:
			res_addr, res_port = get_socket_info( name, port=port,
				protocol=protocol, family=opts.addr_family, socktype=opts.sock_type,
				force_unique_address=opts.force_unique )
		except AddressError:
			if not opts.skip_unresolved: raise
			continue

		# Replace
		src = src[:match.start(group) + src_diff]\
			+ res_addr + src[match.end(group) + src_diff:]
		src_diff += len(res_addr) - len(match.group(group))
		if port:
			src = src[:match.start('port') + src_diff]\
				+ res_port + src[match.end('port') + src_diff:]
			src_diff += len(res_port) - len(match.group('port'))

	(sys.stdout if not opts.out else open(opts.out, 'w')).write(src)

if __name__ == '__main__': sys.exit(main())