bh.py

bh.py

from __future__ import print_function
import os
import sys
from netmiko import ConnectHandler

target_mac = os.environ['TARGET_MAC']
router_ip = os.environ['ROUTER_IP']
router_user = os.environ['ROUTER_USER']
password = os.environ['ROUTER_PW']

home = {
    'device_type': 'cisco_ios',
    'ip':   router_ip,
    'username': router_user,
    'password': password,
    'port' : 22,
    'verbose': False,       
}


def arp_lookup(mac, c):
    show_arp = c.send_command('show arp | i %s' % target_mac)

    try:
        target_ip = show_arp.split()[1]
    except:
        print("target MAC not found")
        sys.exit()

    print("target is %s" % target_ip)
    return target_ip

def execute(cmd_set, c):
    c.enable()
    return c.send_config_set(cmd_set)

def black_hole(mac, c):
    target_ip = arp_lookup(mac, c)
    null_route = [ 'ip route %s 255.255.255.255 null0' % target_ip ]

    return execute(null_route, c)

def remove_black_hole(mac, c):
    target_ip = arp_lookup(mac, c)
    remove_null_route = [ 'no ip route %s 255.255.255.255 null0' % target_ip ]

    return execute(remove_null_route, c)

def main():
    c = ConnectHandler(**home)

    if sys.argv[1] == 'block':
        print(black_hole(target_mac, c))
    else:
        print(remove_black_hole(target_mac, c))

if __name__ == '__main__':
    main()