F5 Pool Member Drain

#!/usr/bin/env python3
"""
F5 Pool Member Drain

READ THE DOCUMENTATION BEFORE YOU RUN THIS:
    https://permitany.com/scripts/f5-pool-drain

It covers what the script does, whether it is safe to run, the
inventory file and credentials it needs, how to read the output,
and what to do when something fails.

    pip install "requests>=2.28" "PyYAML>=6.0"
"""

import argparse
import os
import sys
import time

import requests
import urllib3
import yaml

sys.stdout.reconfigure(encoding="utf-8", errors="replace")


# --- iControl REST client -------------------------------------------------
#
# Inlined verbatim rather than imported: these scripts download standalone
# from the site, so there is no shared module to import. The same block
# appears in every F5 script here; change it in all of them or in none.

def credentials():
    user = os.environ.get("F5_USER")
    password = os.environ.get("F5_PASS")
    if not user or not password:
        raise SystemExit(
            "ERROR: F5_USER and F5_PASS must be set in the environment.\n"
            "\n"
            "  Linux / macOS:   export F5_USER=svc_automation\n"
            "                   export F5_PASS='...'\n"
            "  PowerShell:      $env:F5_USER='svc_automation'\n"
            "                   $env:F5_PASS='...'\n"
            "\n"
            "They are read from the environment on purpose: an argument is\n"
            "visible in `ps auxww` while the process runs and lands in shell\n"
            "history."
        )
    return user, password


class BigIP:
    """Minimal iControl REST client.

    Token auth rather than Basic on every call: a token is scoped, expires,
    and does not put the password on the wire repeatedly.
    """

    def __init__(self, host, user, password, verify=True, timeout=30):
        self.base = "https://%s/mgmt" % host
        self.host = host
        self.timeout = timeout
        self.s = requests.Session()
        self.s.verify = verify
        if not verify:
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        self._login(user, password)

    def _login(self, user, password):
        r = self.s.post(
            self.base + "/shared/authn/login",
            json={"username": user, "password": password, "loginProviderName": "tmos"},
            timeout=self.timeout,
        )
        if r.status_code == 401:
            raise SystemExit("ERROR: %s rejected the credentials (401)." % self.host)
        r.raise_for_status()
        token = r.json()["token"]["token"]
        self.s.headers.update({"X-F5-Auth-Token": token})

    def get(self, path, **kw):
        r = self.s.get(self.base + path, timeout=self.timeout, **kw)
        r.raise_for_status()
        return r.json()

    def patch(self, path, body):
        r = self.s.patch(self.base + path, json=body, timeout=self.timeout)
        r.raise_for_status()
        return r.json()

    def post(self, path, body):
        r = self.s.post(self.base + path, json=body, timeout=self.timeout)
        r.raise_for_status()
        return r.json()


def load_inventory(path):
    """Read the device inventory. Credentials never live in it."""
    try:
        with open(path, "r", encoding="utf-8") as fh:
            doc = yaml.safe_load(fh) or {}
    except FileNotFoundError:
        raise SystemExit("ERROR: inventory file not found: %s" % path)
    except yaml.YAMLError as exc:
        raise SystemExit("ERROR: %s is not valid YAML: %s" % (path, exc))

    defaults = doc.get("defaults") or {}
    devices = doc.get("devices")
    if not devices:
        raise SystemExit(
            "ERROR: %s has no `devices:` list.\n"
            "\n"
            "  defaults:\n"
            "    verify_ssl: true\n"
            "  devices:\n"
            "    - host: ltm01.example.com\n"
            "    - host: 10.0.0.10\n"
            "      verify_ssl: false" % path
        )
    out = []
    for i, d in enumerate(devices, 1):
        merged = dict(defaults)
        merged.update(d or {})
        if not merged.get("host"):
            raise SystemExit("ERROR: device #%d in %s has no `host`." % (i, path))
        out.append(merged)
    return out


def encode(name):
    """tmsh paths use ~ for / in the REST URI: /Common/pool -> ~Common~pool"""
    return name.replace("/", "~")


# --- drain ----------------------------------------------------------------

def member_uri(pool, member):
    return "/tm/ltm/pool/%s/members/%s" % (encode(pool), encode(member))


def current_connections(bip, pool, member):
    stats = bip.get(member_uri(pool, member) + "/stats")
    entries = stats.get("entries", {})
    for v in entries.values():
        nested = v.get("nestedStats", {}).get("entries", {})
        cur = nested.get("serverside.curConns", {}).get("value")
        if cur is not None:
            return int(cur)
    return None


def drain_member(bip, pool, member, wait, poll, force, dry_run):
    """Disable a member and wait for its connections to reach zero.

    The point of the script. Setting a member down and immediately taking the
    server out drops every connection it was serving; the correct sequence is
    disable, then WAIT, then act. Almost nobody waits.

    `session: user-disabled` stops NEW connections but lets existing ones
    finish. `state: user-down` kills them. This uses the first.
    """
    before = current_connections(bip, pool, member)
    print("  %-28s connections=%s" % (member, "unknown" if before is None else before))

    if dry_run:
        print("  [dry-run] would set session=user-disabled and wait up to %ds" % wait)
        return "dry-run"

    bip.patch(member_uri(pool, member), {"session": "user-disabled"})
    print("  disabled (no new connections). Waiting up to %ds..." % wait)

    deadline = time.time() + wait
    last = before
    while time.time() < deadline:
        time.sleep(poll)
        cur = current_connections(bip, pool, member)
        if cur != last:
            print("    connections=%s" % cur)
            last = cur
        if cur == 0:
            print("  DRAINED: 0 connections. Safe to work on the server.")
            return "drained"

    print("  TIMEOUT: still %s connection(s) after %ds." % (last, wait))
    if force:
        bip.patch(member_uri(pool, member), {"state": "user-down"})
        print("  FORCED DOWN as requested. Those connections were reset.")
        return "forced"
    print("  Left disabled, NOT forced down. Existing connections are still being")
    print("  served. Re-run with --force only if resetting them is acceptable.")
    return "timeout"


def enable_member(bip, pool, member, dry_run):
    if dry_run:
        print("  [dry-run] would set session=user-enabled and state=user-up")
        return "dry-run"
    bip.patch(member_uri(pool, member), {"session": "user-enabled", "state": "user-up"})
    print("  %-28s enabled" % member)
    return "enabled"


def resolve_members(bip, pool, wanted):
    data = bip.get("/tm/ltm/pool/%s/members" % encode(pool))
    names = [m["name"] for m in data.get("items", [])]
    if not names:
        raise SystemExit("ERROR: pool %s has no members, or does not exist." % pool)
    if not wanted:
        return names
    missing = [w for w in wanted if w not in names]
    if missing:
        raise SystemExit(
            "ERROR: not a member of %s: %s\nMembers are: %s"
            % (pool, ", ".join(missing), ", ".join(names))
        )
    return wanted


def main():
    ap = argparse.ArgumentParser(
        description="Drain F5 pool members: disable, wait for connections to "
                    "reach zero, and report. Never resets a connection unless "
                    "--force is given."
    )
    ap.add_argument("--inventory", required=True, help="YAML inventory of BIG-IPs")
    ap.add_argument("--pool", required=True, help="Pool name, e.g. /Common/pool_app")
    ap.add_argument("--member", action="append", default=[],
                    help="Member as addr:port. Repeatable. Omit for every member "
                         "(refused unless --all-members is also given)")
    ap.add_argument("--all-members", action="store_true",
                    help="Required to act on every member of the pool")
    ap.add_argument("--enable", action="store_true",
                    help="Re-enable instead of draining")
    ap.add_argument("--wait", type=int, default=300, help="Seconds to wait (default 300)")
    ap.add_argument("--poll", type=int, default=5, help="Seconds between checks (default 5)")
    ap.add_argument("--force", action="store_true",
                    help="After the wait, force the member down, RESETTING any "
                         "remaining connections")
    ap.add_argument("--dry-run", action="store_true",
                    help="Show what would happen and change nothing")
    args = ap.parse_args()

    if not args.member and not args.all_members:
        raise SystemExit(
            "ERROR: specify --member addr:port, or pass --all-members to act on\n"
            "the whole pool. Draining every member at once takes the service\n"
            "offline, so it has to be asked for explicitly."
        )

    user, password = credentials()
    devices = load_inventory(args.inventory)
    worst = 0

    for dev in devices:
        host = dev["host"]
        print("=" * 72)
        print("%s  pool=%s" % (host, args.pool))
        print("=" * 72)
        try:
            bip = BigIP(host, user, password,
                        verify=dev.get("verify_ssl", True),
                        timeout=dev.get("timeout", 30))
            members = resolve_members(bip, args.pool, args.member)

            if args.all_members and not args.member and not args.enable:
                print("WARNING: draining ALL %d members. The pool will have no"
                      % len(members))
                print("         available member when this finishes.")

            for m in members:
                if args.enable:
                    enable_member(bip, args.pool, m, args.dry_run)
                else:
                    result = drain_member(bip, args.pool, m, args.wait,
                                          args.poll, args.force, args.dry_run)
                    if result == "timeout":
                        worst = max(worst, 1)
        except requests.exceptions.SSLError as exc:
            print("ERROR: TLS failed for %s: %s" % (host, exc))
            print("       Set verify_ssl: false in the inventory only if the")
            print("       certificate is a known self-signed one.")
            worst = 2
        except requests.exceptions.RequestException as exc:
            print("ERROR: %s unreachable: %s" % (host, exc))
            worst = 2
        except SystemExit:
            raise
        print()

    return worst


if __name__ == "__main__":
    sys.exit(main())