#!/usr/bin/python
# Copyright 2016 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Apply ip address and default route for a set of network interfaces."""

import errno
import json
import os
import re
import socket
import struct
import subprocess
import sys
import options

optspec = """
ipapply <interface names>
--
"""

# TODO(poist): Add support for alias and mtu.

# Allow test override of configuration file location.
base_path = os.getenv('IPAPPLY_CONFIG_FILE_BASE', '')

STATIC_IP_CONFIGS_DIR = base_path + '/config/ip/static'
DYNAMIC_IP_CONFIGS_DIR = base_path + '/tmp/ip/dynamic'
TMP_CONMAN_DIR = base_path + '/tmp/conman'


def AtomicWrite(filename, data):
  tmp_filename = filename + '.tmp'
  with open(tmp_filename, 'w') as tmp:
    tmp.write(data)
  os.rename(tmp_filename, filename)


def Log(s, *args):
  sys.stdout.flush()
  ss = 'ipapply: %s\n' % (s,)
  if args:
    sys.stderr.write(ss % args)
  else:
    sys.stderr.write(ss)
  sys.stderr.flush()


def LoadJsonFromFile(path):
  """Deserializes a JSON file to a Python object.

  Args:
    path: The path to the JSON file to be converted.

  Returns:
    Dict containing the JSON data.

  Raises:
    OSError/IOError: If the JSON data file could not be loaded.
    ValueError: If the JSON data loaded from the file does not contain the
                requested key, a ValueError is raised.
  """
  if not os.path.exists(path):
    # It is a valid state that there might not be a JSON data file yet,
    # so just return without printing an error.
    return

  with open(path) as f:
    return json.load(f)


def InterfaceExists(iface_name):
  """Return True if the named network interface exists."""
  return os.path.isdir(os.path.join(base_path + '/sys/class/net', iface_name))


def IsInterfaceUp(iface_name):
  """Return True if the named network interface is up."""
  if subprocess.call(['is-interface-up', iface_name]) == 0:
    return True
  return False


def IsProcessRunning(pid):
  """Check whether a process is running."""
  pid = int(pid)
  if pid <= 0:
    raise ValueError('invalid pid')

  try:
    os.kill(pid, 0)
  except OSError as err:
    if err.errno == errno.ESRCH:
      # ESRCH == No such process
      return False
    elif err.errno == errno.EPERM:
      # EPERM clearly means there's a process to deny access to
      return True
    else:
      # According to "man 2 kill" possible error values are
      # (EINVAL, EPERM, ESRCH)
      raise


def IpAddrUpdate(iface_name, new_ips):
  """Add new address (with /netmask) on an interface, and remove old ones."""
  ip_show = subprocess.check_output(['ip', 'addr', 'show', 'dev', iface_name])
  current_ips = set()

  pattern = r'\s*inet\s+(?P<ip>(\d{1,3}\.?){4})(?P<netmask>\/\d{1,2})?'
  regex = re.compile(pattern)
  for line in ip_show.splitlines():
    match = re.match(pattern, line)
    if match:
      current_ips.add(match.group('ip') + (match.group('netmask') or '/32'))

  for new_ip in new_ips - current_ips:
    IpAddrAdd(iface_name, new_ip)

  for old_ip in current_ips - new_ips:
    IpAddrDel(iface_name, old_ip)


def IpAddrAdd(iface, ip):
  Log('Adding address %r on interface %r', ip, iface)
  subprocess.check_call(['ip', 'addr', 'add', ip, 'dev', iface])
  # We don't want 'proto kernel' routes.
  ip_route_show = ['ip', 'route', 'show', 'dev', iface]
  for line in subprocess.check_output(ip_route_show).splitlines():
    if 'proto kernel' in line:
      route = line.split()[0]
      if subprocess.call(['ip', 'route', 'del', route, 'dev', iface]) != 0:
        Log('Failed to delete route %s on interface %s', route, iface)


def IpAddrDel(iface_name, ip):
  Log('Deleting address %r on interface %r', ip, iface_name)
  subprocess.check_call(['ip', 'addr', 'del', ip, 'dev', iface_name])


def RouteAdd(iface_name, routers=None, subnet=None):
  """Set a default route via this interface through apman."""
  # apman supports only a single route, and installs it as 'default'.
  # arbitrarily select the first route if multiple provided.
  subnet_path = os.path.join(TMP_CONMAN_DIR, 'subnet.%s' % iface_name)
  gateway_path = os.path.join(TMP_CONMAN_DIR, 'gateway.%s' % iface_name)

  if subnet:
    Log('add subnet route: %r on dev %r', subnet, iface_name)
    AtomicWrite(subnet_path, subnet)

  if routers:
    first_route = routers.split()[0]
    Log('add default route on dev %r with routers: %r',
        iface_name, first_route)
    AtomicWrite(gateway_path, first_route)


def SelectAndApply(iface_name):
  """Select the config file and apply the contents."""
  static_config_file = os.path.join(STATIC_IP_CONFIGS_DIR, iface_name)
  dynamic_config_file = os.path.join(DYNAMIC_IP_CONFIGS_DIR, iface_name)

  # Select the appropriate configuration file.
  # Normally static configuration takes precedence over dynamic
  # configuration. If this is a box that may be running in bridge
  # mode (ie with both a wan [uplink] port, and a lan port), decide
  # which configuration based on the current state of the wan port.
  selected_config_file = None
  if iface_name == 'br0' and InterfaceExists('wan0'):
    # Device that supports bridge mode, select based on wan link status.
    if IsInterfaceUp('wan0'):
      selected_config_file = static_config_file
    else:
      selected_config_file = dynamic_config_file
  elif IsInterfaceUp(iface_name):
    if os.path.isfile(static_config_file):
      selected_config_file = static_config_file
    else:
      selected_config_file = dynamic_config_file
  else:
    Log('interface %r down, clearing config', iface_name)

  if not selected_config_file or not os.path.isfile(selected_config_file):
    Log('no available configuration file for interface %r', iface_name)
    config = {}
  else:
    config = LoadJsonFromFile(selected_config_file)
    if config:
      config = config[iface_name]
    else:
      Log('failed to load JSON configuration from %r for interface %r',
          selected_config_file, iface_name)

  new_ips = set()
  for ip in config.get('ip', []):
    if 'new_ip_address' in ip:
      ip_string = ip['new_ip_address']
      netmask = '32'
      if 'new_subnet_mask' in ip:
        netmask = str(NormalizeNetmask(ip['new_subnet_mask']))
      ip_string = ip_string + '/' + netmask
      new_ips.add(ip_string)

  IpAddrUpdate(iface_name, new_ips)

  for ip in config.get('ip', []):
    subnet = None
    if 'new_ip_address' in ip and 'new_subnet_mask' in ip:
      mask = str(NormalizeNetmask(ip['new_subnet_mask']))
      subnet = GetSubnet(ip['new_ip_address'], mask)
    RouteAdd(iface_name, routers=ip.get('new_routers', None), subnet=subnet)


def IpToInt(ip):
  return struct.unpack('!I', socket.inet_pton(socket.AF_INET, ip))[0]


def IntToIp(ip):
  return socket.inet_ntop(socket.AF_INET, struct.pack('!I', ip))


def NormalizeNetmask(netmask):
  """Convert a netmask to 'bits', e.g. 24, if it isn't already."""
  if not isinstance(netmask, basestring):
    raise ValueError('NormalizeNetmask expects a string')

  if '.' in netmask:
    return NetmaskToBits(netmask)

  # If we got an invalid string, i.e. one that isn't an IP string (handled
  # above) or an integer, it will be detected here and a ValueError will be
  # raised.
  x = int(netmask)
  if not 0 <= x <= 32:
    raise ValueError('Netmask bits should be between 0 and 32')

  return x


def NetmaskToBits(netmask):
  """255.255.255.0' -> 24"""

  # '255.255.255.0' -> '11111111111111111111111100000000'
  netmask_bitstring = bin(IpToInt(netmask))[2:]

  # '11111111111111111111111100000000' -> 24
  return len(netmask_bitstring.split('0')[0])


def ApplyNetmask(ip, netmask):
  int_ip = IpToInt(ip)
  netmask = int(NormalizeNetmask(netmask))
  int_netmask = (2**netmask - 1) << (32 - netmask)
  return IntToIp(int_ip & int_netmask)


def GetSubnet(ip, netmask):
  return ''.join((ApplyNetmask(ip, netmask), '/', netmask))


def main():
  o = options.Options(optspec)
  (_, _, extra) = o.parse(sys.argv[1:])
  if not extra:
    o.fatal('at least one interface name expected')

  # if using NFS root, don't mess with interfaces other than
  # wireless, or we could break the NFS connection.
  is_nfs = os.path.exists('/tmp/NFS')

  # uniquify the set of interface names
  for iface_name in set(extra):
    if iface_name.startswith('wcli') or not is_nfs:
      SelectAndApply(iface_name)


if __name__ == '__main__':
  main()