ip-ext

#!/usr/bin/env python3

import hashlib as hl, ipaddress as ip, subprocess as sp, itertools as it
import os, sys, logging, re, json


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


def ipv6_dns(opts):
	addr_ex = ip.IPv6Address(opts.ipv6_address).exploded
	log.debug('Exploded address: {}', addr_ex)
	addr_digits = addr_ex.replace(':', '')
	if opts.exploded: print(addr_ex)
	if opts.djbdns: print(addr_digits)
	if not (opts.djbdns or opts.exploded): print('.'.join(reversed(addr_digits)) + '.ip6.arpa')


def ipv6_lladdr(opts):
	if opts.assign and not opts.iface:
		parser.error('--assign requires --iface option to be specified.')
	addr_print = not opts.assign

	ipv6_norm = lambda a: ( '' if not a else
		ip.IPv6Address(a).compressed.lower() )
	def mac_norm(a):
		if not a: return ''
		a = ''.join(('00' + v)[-2:] for v in a.split(':'))
		return ':'.join(a[n:n+2] for n in range(0, 6*2, 2)).lower()

	if_mac = if_lladdr = None
	if opts.iface:
		iface_info = json.loads(sp.run(
			'ip --json addr'.split(), check=True, stdout=sp.PIPE ).stdout)
		for iface in iface_info:
			if iface['ifname'] != opts.iface: continue
			if_mac = iface.get('address')
			for a in iface.get('addr_info') or list():
				if not (a.get('family') == 'inet6' and a.get('scope') == 'link'): continue
				if_lladdr = a.get('local')
			break
		else:
			opts.parser.error( 'Specified --iface name not found:'
				' {} [{}]'.format(opts.iface, ', '.join(v['ifname'] for v in iface_info)) )
	if_mac, if_lladdr = mac_norm(if_mac), ipv6_norm(if_lladdr)

	mac = lladdr = None
	if opts.lladdr and opts.mac:
		opts.parser.error('Only one of --lladdr/--mac can be specified, not both')
	elif opts.mac: mac = mac_norm(opts.mac)
	elif opts.lladdr: lladdr = ipv6_norm(opts.lladdr)
	else:
		mac, lladdr = if_mac, if_lladdr
		if lladdr and opts.assign: opts.assign = False # just run sanity-checks

	if mac:
		lladdr_chk = lladdr
		parts = mac.split(':')
		parts.insert(3, 'ff')
		parts.insert(4, 'fe')
		parts[0] = f'{int(parts[0], 16) ^ 2:x}'
		lladdr = ipv6_norm('fe80::' + ':'.join(
			''.join(parts[n:n+2]) for n in range(0, len(parts), 2) ))
		if lladdr_chk and lladdr != lladdr_chk:
			log.error('mac/lladdr mismatch: {} -> {} [not {}]', mac, lladdr, lladdr_chk)
			return 1

	if lladdr:
		mac_chk = mac
		parts = list(it.chain.from_iterable( (p[:2], p[-2:])
			for p in (('0000' + p)[-4:] for p in lladdr.split(':')[-4:]) ))
		parts[0] = f'{int(parts[0], 16) ^ 2:02x}'
		parts[3:5] = []
		mac = mac_norm(':'.join(parts))
		if mac_chk and mac != mac_chk:
			log.error('mac/lladdr mismatch: {} -> {} [not {}]', lladdr, mac, mac_chk)
			return 1

	if addr_print:
		if opts.iface: print('iface: ', opts.iface)
		print('mac:   ', mac)
		print('lladdr:', lladdr)
	if opts.assign:
		if if_lladdr == lladdr:
			log.debug( 'Specified link-local address'
				' already assigned to iface [{}]: {}', opts.iface, lladdr )
		if sp.Popen([ 'ip', 'addr', 'add',
				f'{lladdr}/64', 'dev', opts.iface, 'scope', 'link' ]).wait():
			log.error( 'Failed to assign link-local'
				' address ({}) to interface: {}', lladdr, opts.iface)
			return 1
		log.debug('Assigned lladdr to interface [{}]: {}', opts.iface, lladdr)


ipv6_name_b2_person = 'ipx'
def ipv6_name(opts):
	pre = ip.IPv6Network(opts.prefix or '::/0', strict=False)
	if pre.prefixlen == 128:
		pre_addr = int.from_bytes(pre.network_address.packed, 'big')
		for n in range(129):
			if pre_addr & 1 << n: break
		pre = ip.IPv6Network(pre.network_address.compressed + f'/{128 - n}')
		log.debug('Auto-detected prefix-length for address: {}', pre)
	suffix, suffix_bits = opts.name.encode('utf-8'), 128 - pre.prefixlen

	b2_params, opt = ['+', '', ipv6_name_b2_person], opts.blake2_params
	if opt.find(':') != opt.find('::'): opt = opt.replace('::', '\ue000')
	opt = list(s.replace('\ue000', ':') for s in opt.split(':'))
	b2_params[:len(opt)] = opt
	b2_len, b2_key, b2_person = b2_params
	b2_len = suffix_bits if b2_len == '+' else int(b2_len) * 8
	if b2_len > suffix_bits:
		opts.parser.error('Suffix length in --blake2-params is too long for specified -p/--prefix')
	b2_person, b2_key = b2_person.encode(), b2_key.encode()
	if opts.blake2:
		suffix = int.from_bytes(hl.blake2s( suffix,
			digest_size=16, key=b2_key, person=b2_person ).digest(), 'big')
		suffix &= (2**b2_len - 1)
	else:
		if len(suffix) * 8 > suffix_bits:
			opts.parser.error('Encoded name length is too long for specified -p/--prefix')
		suffix = int.from_bytes(suffix, 'big')
	addr = ip.IPv6Address(int.from_bytes(pre.network_address.packed, 'big') | suffix)
	log.debug( 'Generated address: {}'
		' [prefix-bits={} suffix-bits={}]', addr, pre.prefixlen, suffix_bits )

	addr = addr.compressed if not opts.exploded else addr.exploded
	if opts.mask: addr += f'/{pre.prefixlen}'
	print(addr)


def iptables_flush(opts):
	for ipt in 'iptables', 'ip6tables':
		tables = sp.check_output(f'{ipt}-save')
		tables_restore = sp.Popen(f'{ipt}-restore', stdin=sp.PIPE)
		tables_restore_line = ( lambda line_new=None:
			tables_restore.stdin.write((line_new or line).encode() + b'\n') )
		for line in tables.decode().splitlines():
			if line.startswith('*') or line == 'COMMIT': tables_restore_line()
			if re.search('^:[A-Z]+', line): tables_restore_line(f'{line.split()[0]} ACCEPT')
		tables_restore.stdin.close()
		if tables_restore.wait() != 0:
			raise RuntimeError(f'{ipt}-restore exited with error status')


def main(args=None):
	import argparse
	parser = argparse.ArgumentParser(
		description='Tool to manipulate IP addresses, network interfaces and misc settings.')
	parser.add_argument('-d', '--debug', action='store_true', help='Verbose operation mode.')

	cmds = parser.add_subparsers(title='Supported actions', dest='call')

	cmd = cmds.add_parser('ipv6-lladdr',
		help='Generate and/or assign link-local IPv6 address to the interface.')
	cmd.add_argument('-i', '--iface', metavar='dev',
		help='Interface name to print mac/ipv6-lladdr for.')
	cmd.add_argument('-x', '--assign', action='store_true',
		help='Assign address to the interface,'
			' unless already there. Requires -i/--iface option to be used.')
	cmd.add_argument('-l', '--lladdr', metavar='ipv6',
		help='Link-local IPv6 address to convert to mac address.'
			' Can be used with -x/--assign to set this arbitrary lladdr after a sanity-check.')
	cmd.add_argument('-m', '--mac', metavar='mac',
		help='MAC address to convert to link-local address.'
			' Can be used with -x/--assign to set lladdr from arbitrary mac.')

	cmd = cmds.add_parser('ipv6-dns',
		help='Convert IPv6 address to "*.ip6.arpa" (or other) domain name format.')
	cmd.add_argument('ipv6_address', help='Address to convert.')
	cmd.add_argument('-e', '--exploded',
		action='store_true', help='Convert to exploded IPv6 address for an SPF record.')
	cmd.add_argument('-d', '--djbdns',
		action='store_true', help='Convert to canonical djbdns record format instead.')

	cmd = cmds.add_parser('ipv6-name',
		help='Convert non-secret name to a consistent IPv6 address in a number ways.')
	cmd.add_argument('name', help='Name to convert.')
	cmd.add_argument('-p', '--prefix', metavar='prefix',
		help='IPv6 prefix (or network) to combine generated suffix with.'
			' Exits with error if generated suffix does not fit into address with prefix.')
	cmd.add_argument('-e', '--exploded', action='store_true',
		help='Print resulting address in an exploded/padded format.')
	cmd.add_argument('-m', '--mask', action='store_true',
		help='Print netmask from prefix in the resulting address.')
	cmd.add_argument('-u', '--utf8', action='store_true',
		help='Encode to utf8 and use that as suffix. Default mode.')
	cmd.add_argument('-b', '--blake2', action='store_true',
		help='Use blake2s hash after utf8-encoding with specified parameters.')
	cmd.add_argument('--blake2-params',
		metavar='byte-len(:key(:person))', default=f'+::{ipv6_name_b2_person}',
		help='Byte-length/person/key parameters for blake2s hash with -b/--blake2 option.'
			' Special "+" value for length will use all bytes outside of specified prefix - up to 16.'
			' Key/person are utf-8 only, ":" should be escaped as "::". Default: %(default)s')

	cmd = cmds.add_parser('iptables-flush',
		help='Flush all iptables rules (all chains/tables, ipv4/ipv6) using CLI save/restore tools.')

	opts = parser.parse_args(sys.argv[1:] if args is None else args)
	opts.parser = parser

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

	try:
		func = dict( ipv6_dns=ipv6_dns, ipv6_lladdr=ipv6_lladdr,
			ipv6_name=ipv6_name, iptables_flush=iptables_flush )[opts.call.replace('-', '_')]
	except KeyError:
		parser.error('Action {!r} is not implemented.'.format(opts.call))
	return func(opts)

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