common.py

#!/usr/bin/env python3
#
# Copyright (c) 2017,2018, F5 Networks, Inc.
#
# 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.

"""Common utility functions."""

import ipaddress
import re
import sys
import time
import json
import logging
import socket
import argparse

import jwt
import requests

from requests.auth import AuthBase


# Big-IP Address Pattern: <ipaddr>%<route_domain>
ip_rd_re = re.compile(r'^([^%]*)%(\d+)$')


def parse_log_level(log_level_arg):
    """Parse the log level from the args.

    Args:
        log_level_arg: String representation of log level
    """
    LOG_LEVELS = ['CRITICAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG']
    if log_level_arg not in LOG_LEVELS:
        msg = 'Invalid option: {0} (Valid choices are {1})'.format(
            log_level_arg, LOG_LEVELS)
        raise argparse.ArgumentTypeError(msg)

    log_level = getattr(logging, log_level_arg, logging.INFO)

    return log_level


def setup_logging(logger, log_format, log_level):
    """Configure logging."""
    logger.setLevel(log_level)

    formatter = logging.Formatter(log_format)

    consoleHandler = logging.StreamHandler(sys.stdout)
    consoleHandler.setFormatter(formatter)
    logger.addHandler(consoleHandler)
    logger.propagate = False


def set_marathon_auth_args(parser):
    """Set the authorization for Marathon."""
    parser.add_argument("--marathon-auth-credential-file",
                        env_var='F5_CC_MARATHON_AUTH',
                        help="Path to file containing a user/pass for "
                        "the Marathon HTTP API in the format of 'user:pass'.")
    parser.add_argument("--dcos-auth-credentials",
                        env_var='F5_CC_DCOS_AUTH_CREDENTIALS',
                        help="DC/OS service account credentials")
    parser.add_argument("--dcos-auth-token",
                        env_var='F5_CC_DCOS_AUTH_TOKEN',
                        help="DC/OS ACS Token")

    return parser


class DCOSAuth(AuthBase):
    """DCOSAuth class.

    Manage authorization credentials for DCOS
    """

    def __init__(self, credentials, ca_cert, token):
        """Initialize DCOSAuth."""
        if credentials:
            creds = json.loads(credentials)
            self.scheme = creds['scheme']
            self.uid = creds['uid']
            self.private_key = creds['private_key']
            self.login_endpoint = creds['login_endpoint']
        self.token = token
        self.verify = False
        self.auth_header = None
        self.expiry = 0
        if ca_cert:
            self.verify = ca_cert

    def __call__(self, auth_request):
        """Get the ACS token."""
        if self.token:
            self.auth_header = 'token=' + self.token
            auth_request.headers['Authorization'] = self.auth_header
            return auth_request

        if not self.auth_header or int(time.time()) >= self.expiry - 10:
            self.expiry = int(time.time()) + 3600
            payload = {
                'uid': self.uid,
                # This is the expiry of the auth request params
                'exp': int(time.time()) + 60,
            }
            token = jwt.encode(payload, self.private_key, self.scheme)

            data = {
                'uid': self.uid,
                'token': token.decode('ascii'),
                # This is the expiry for the token itself
                'exp': self.expiry,
            }
            r = requests.post(self.login_endpoint,
                              json=data,
                              timeout=(3.05, 46),
                              verify=self.verify)
            r.raise_for_status()

            self.auth_header = 'token=' + r.cookies['dcos-acs-auth-cookie']

        auth_request.headers['Authorization'] = self.auth_header
        return auth_request


def get_marathon_auth_params(args):
    """Get the Marathon credentials."""
    marathon_auth = None
    if args.marathon_auth_credential_file:
        with open(args.marathon_auth_credential_file, 'r') as f:
            line = f.readline().rstrip('\r\n')

        if line:
            marathon_auth = tuple(line.split(':'))
    elif args.dcos_auth_credentials or args.dcos_auth_token:
        return DCOSAuth(args.dcos_auth_credentials, args.marathon_ca_cert,
                        args.dcos_auth_token)

    if marathon_auth and len(marathon_auth) != 2:
        print(
            "Please provide marathon credentials in user:pass format"
        )
        sys.exit(1)

    return marathon_auth


def set_logging_args(parser):
    """Add logging-related args to the parser."""
    parser.add_argument("--log-format",
                        env_var='F5_CC_LOG_FORMAT',
                        help="Set log message format",
                        default="%(asctime)s %(name)s: %(levelname)"
                        " -8s: %(message)s")
    parser.add_argument("--log-level",
                        env_var='F5_CC_LOG_LEVEL',
                        type=parse_log_level,
                        help="Set logging level. Valid log levels are: "
                        "DEBUG, INFO, WARNING, ERROR, and CRITICAL",
                        default='INFO')
    return parser


ip_cache = dict()


def resolve_ip(host):
    """Get the IP address for a hostname."""
    cached_ip = ip_cache.get(host, None)
    if cached_ip:
        return cached_ip
    else:
        try:
            ip = socket.gethostbyname(host)
            ip_cache[host] = ip
            return ip
        except socket.gaierror:
            return None


def split_ip_with_route_domain(address):
    u"""Return ip and route-domain parts of address

    Input ip format must be of the form:
        <ip_v4_or_v6_addr>[%<route_domain_id>]
    """
    match = ip_rd_re.match(address)
    if match:
        ip = match.group(1)
        route_domain = int(match.group(2))
    else:
        ip = address
        route_domain = None

    return ip, route_domain


def validate_bigip_address(address):
    """Verify the address is a valid Big-IP address"""

    is_valid = True
    try:
        ip = split_ip_with_route_domain(address)[0]
        ipaddress.ip_address(ip)
    except Exception:
        is_valid = False
    return is_valid