#!/usr/bin/env python2.7
# -*- coding: utf-8 -*-
#
# Copyright © 2013 Zulip, Inc.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

# Original author: Luke Faraone

'''Configure a host for EC2-VPC dynamically assigned network interfaces

Amazon VPC gives us a good deal of flexibility compared to classic EC2.
However there are limitations; you can assign multiple IPs to a host
yet only the first IP per interface will be DHCP assigned, and you are
limited in the total number of interfaces you have, so doing one-IP-per-
interface is also untenable.

This script grabs the metadata provided by AWS and uses it to correctly
configure all available network interfaces.

It is suitable to be hooked in to system boot and network
reconfiguration scripts.

Note that it currently does not handle the deconfiguration of
interfaces.

'''

from __future__ import absolute_import
import sys
import logging
import logging.handlers
import subprocess

import boto.utils
import netifaces
if False:
    # See https://zulip.readthedocs.io/en/latest/testing/mypy.html#mypy-in-production-scripts
    from typing import Optional

def address_of(device_id):
    # type: (int) -> Optional[str]
    try:
        return netifaces.ifaddresses("ens%i" % (device_id,))[netifaces.AF_INET][0]['addr']
    except KeyError:
        return None

def guess_gateway(device_id):
    # type: (int) -> Optional[str]
    # This will not work if the default gateway isn't n.n.n.1.
    address = address_of(device_id)
    if address is None:
        return None
    gateway = address.split('.')
    gateway[3] = '1'
    return '.'.join(gateway)

log = logging.getLogger('configure-cloud-interfaces')
log.setLevel(logging.DEBUG)

log.addHandler(logging.handlers.SysLogHandler(facility=logging.handlers.SysLogHandler.LOG_DAEMON))
log.addHandler(logging.StreamHandler())
log.info("Starting.")

metadata = boto.utils.get_instance_metadata()
if metadata is None:
    log.error("Could not get instance metadata!")
    sys.exit(1)

macs = metadata["network"]["interfaces"]["macs"]
ids = [int(macdata['device-number']) for macdata in macs.values()]
ifaces = [iface for iface in netifaces.interfaces() if ":" not in iface and iface != "lo"]

# Number of IDs should equal number of interfaces
if len(ids) != len(ifaces):
    log.error("Metadata indicated %i interfaces but we have %i!" % (len(ids), len(ifaces)))
    sys.exit(1)

for device in macs.values():
    # There's an annoying API inconsistency here:
    #   If you have multiple IPs, local-ipv4s is a list.
    #   If you only have one, local-ipv4s is a string.
    # Who knew?
    if isinstance(device['local-ipv4s'], str):
        # Only do dhcp, don't try to assign addresses
        to_configure = [device['local-ipv4s']]
    else:
        to_configure = list(device['local-ipv4s'])
    device_number = int(device['device-number'])
    if subprocess.check_output(["lsb_release", '-sc']).strip() == "xenial":
        # HACK: It appears on Xenial device_number in the data set is 0 but in reality it is 3
        device_number += 3

    address = address_of(device_number)

    if address is None:
        # If the device was not autoconfigured, do so now.
        log.info("Device ens%i not configured, starting dhcpd" % (device_number,))
        subprocess.check_call(['/sbin/dhcpcd', 'ens%i' % (device_number,)])

        dev_num = str(device_number)
        address = address_of(device_number)
        gateway = guess_gateway(device_number)
        assert(address is not None)
        assert(gateway is not None)

        # Horrible hack to route return packets on the correct interface
        # See http://unix.stackexchange.com/a/4421/933
        subprocess.check_call(
            ['/sbin/ip', 'rule', 'add', 'fwmark', dev_num, 'table', dev_num])
        subprocess.check_call(
            ['/sbin/ip', 'route', 'add', '0.0.0.0/0', 'table', dev_num, 'dev',
             'ens%i' % (device_number,), 'via', gateway])
        subprocess.check_call(
            ['/sbin/iptables', '-t', 'mangle', '-A', 'OUTPUT', '-m', 'conntrack', '--ctorigdst',
             address, '-j', 'MARK', '--set-mark', dev_num])

    to_configure.remove(address)

    for (count, ip) in enumerate(to_configure):
        # Configure the IP via a virtual interface
        device = "ens%i:%i" % (device_number, count)
        log.info("Configuring %s with IP %s" % (device, ip))
        subprocess.check_call(['/sbin/ifconfig', device, ip])
        subprocess.check_call(
            ['/sbin/iptables', '-t', 'mangle', '-A', 'OUTPUT', '-m', 'conntrack', '--ctorigdst',
             ip, '-j', 'MARK', '--set-mark', str(device_number)])

    for throwaway in range(2):
        # Don't freak out if this doens't work.
        subprocess.call(
            ['/sbin/ip', 'route', 'del', '10.0.0.0/8'])

log.info("Finished.")