example of how to drain/restore traffic on a Juniper device

traffic_management.py

import argparse

import logging
import sys
import time
from jnpr.junos import Device
from jnpr.junos.utils.config import Config

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logger = logging.getLogger(__name__)
requests_logger = logging.getLogger('ncclient')
requests_logger.setLevel(logging.ERROR)


def drain_traffic(device):
    with Device(host=device, user='root', password='Password') as dev:
        logger.info('got facts:\n {}'.format(dev.facts))
        if dev.facts['hostname'] != device:
            raise ValueError('The hostname on the device does not match the'
                             ' hostname we want to drain. got: {} expected: {}'.format(dev.facts['hostname'], device))
        with Config(dev, mode='exclusive') as cu:
            cu.load('set protocols ospf apply-groups DRAIN',
                    format='set')
            logger.info('diff just prior to the commit: {}'.format(cu.diff()))
            cu.commit(comment='Draining traffic', confirm=10)
            logger.info('commit applied, waiting 30 seconds to confirm')
            time.sleep(30)
            cu.commit()
            logger.info('Device drained!')


def restore_traffic(device):
    with Device(host=device, user='root', password='Password') as dev:
        logger.info('got facts:\n {}'.format(dev.facts))
        if dev.facts['hostname'] != device:
            raise ValueError('The hostname on the device does not match the'
                             ' hostname we want to restore. got: {} expected: {}'.format(dev.facts['hostname'], device))
        with Config(dev, mode='exclusive') as cu:
            cu.load('delete protocols ospf apply-groups DRAIN',
                    format='set')
            logger.info('diff just prior to the commit: {}'.format(cu.diff()))
            cu.commit(comment='Restoring traffic', confirm=10)
            logger.info('commit applied, waiting 30 seconds to confirm')
            time.sleep(30)
            cu.commit()
            logger.info('Device restored!!')


# mapping of operations.
mappings = {'drain': drain_traffic,
            'restore': restore_traffic}

parser = argparse.ArgumentParser(description='Manage traffic management operations against Juniper devices.')
parser.add_argument('--device', type=str,
                    help='The name of the device to perform operations on')
parser.add_argument('--operation', type=str,
                    help='The name of the operation to perform on the device. One of {}'.format(mappings.keys()),
                    required=True, choices=mappings.keys())
args = parser.parse_args()


logger.info('Going to begin draining traffic on {}'.format(args.device))
mappings[args.operation](args.device)

context: https://medium.com/@dan.lindow/an-intro-into-juniper-automation-4024ad3055be