# Copyright 2020 Yamaha Corporation
# 
# 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.

import os
import json
import asyncio
from pysnmp.hlapi.v3arch.asyncio import (
    get_cmd, SnmpEngine, CommunityData, UdpTransportTarget,
    ContextData, ObjectType, ObjectIdentity
)
import boto3


async def is_router_and_tunnel_active(host):
    transport = await UdpTransportTarget.create((host, 161), timeout=1, retries=1)
    errorIndication, errorStatus, errorIndex, varBinds = await get_cmd(
        SnmpEngine(),
        CommunityData('public'),
        transport,
        ContextData(),
        ObjectType(ObjectIdentity('1.3.6.1.4.1.1182.2.3.16.1.2.1195'))
    )

    if (not errorIndication) and (not errorStatus) and (varBinds[0][1] == 1):
        return True

    return False


def update_vpc_route(table, cidr, eni):

    boto3.client('ec2').replace_route(
        DestinationCidrBlock=cidr,
        NetworkInterfaceId=eni,
        RouteTableId=table,
    )


def lambda_handler(event, context):

    MASTER_IP = os.environ['MASTER_IP']
    MASTER_ENI = os.environ['MASTER_ENI']
    SLAVE_ENI = os.environ['SLAVE_ENI']

    ROUTE_TABLE = os.environ['ROUTE_TABLE']
    TARGET_CIDR = os.environ['TARGET_CIDR']

    active = asyncio.run(is_router_and_tunnel_active(MASTER_IP))

    update_vpc_route(ROUTE_TABLE, TARGET_CIDR, MASTER_ENI if active else SLAVE_ENI)

    return {
        'isBase64Encoded': False,
        'statusCode': 200 if active else 503,
        'statusDescription': '200 OK' if active else '503 Service Unavailable',
        'headers': {
            'Content-Type': 'application/json'
        },
        'body': json.dumps(active)
    }
