Edit

Remove expired vault tokens

The guide describes how to remove Vault tokens in the ZooKeeper's backend.

The procedure uses two scripts executed on the INT1 server (the host with installed ZooKeeper) consecutively in the correct order.

Prepare environment

To prepare the environment, copy the scripts to the INT1 server:

1. Log in to the INT1 server via SSH.

2. Switch to wf_user.

3. Go to the /tmp directory.

4. Inside the directory, create the following files:

   • zookeeper-cleanup.py

   • regenerate_root_token.sh

5. Paste the following contents to the files:

zookeeper-cleanup.py

regenerate_root_token.sh

#!/usr/bin/env python

import argparse
import logging
import os

from datetime import datetime
from kazoo.client import KazooClient
from kazoo.retry import KazooRetry

VAULT_DEFAULT_CONFIG_FILE_PATH = '/opt/workfusion/vault/server.hcl'
VAULT_DEFAULT_HOST = 'localhost:2181'
ZOOKEEPER_DEFAULT_CLEANUP_PATHS = '/vault/sys/token/id,/vault/sys/expire/id/auth/cert/login'

# Logging
log_default_level = logging.INFO
logger = logging.getLogger('zookeeper-cleanup')
log_formatter = logging.Formatter("[%(asctime)s][%(levelname).4s] %(message)s")
log_handler = logging.StreamHandler()
log_handler.setFormatter(log_formatter)
logger.addHandler(log_handler)
logger.setLevel(log_default_level)


def get_ttl_from_vault_config(config_file):
    """ Returns TTL in seconds read from vault configuration (server.hcl file) """
    VAULT_TTL_CONFIG_NAME = 'max_lease_ttl'

    with open(config_file) as f:
        for line in f:
            if VAULT_TTL_CONFIG_NAME in line:
                ttl = ' '.join(line.split()).replace('"', '').replace(' ', '').split('=')[1] # Removing spaces, quotes and taking number
                break

    return int(ttl)


class ZookeeperCleanupManager():
    def __init__(self, zookeeper_hosts, znodes_to_cleanup, vault_ttl):
        self.znodes_to_cleanup = znodes_to_cleanup.split(',')
        self.vault_ttl = vault_ttl
        self.zk_retry = KazooRetry(max_tries=1000, delay=0.5, backoff=2)
        self.zk_client = KazooClient(hosts=zookeeper_hosts, timeout=60, connection_retry=self.zk_retry, command_retry=self.zk_retry)

    def is_znode_expired(self, znode_data):
        """ Returns True if znode_data's last modification was earlier than self.znode_ttl_seconds from now(). Returns False otherwise. """
        last_modified_timestamp = znode_data.mtime // 1000

        return True if datetime.fromtimestamp(last_modified_timestamp + self.vault_ttl) < datetime.now() else False

    def remove_expired_children(self, znode_path, dry_run=False):
        """ Function returns a dict with count of all, expired and removed children (the ones that are not parents themselves). """
        if not self.zk_client.exists(znode_path):
            logger.warning("%s doesn't exist.", znode_path)
            return {'all': 0, 'expired': 0, 'removed': 0}

        children = self.zk_client.get_children(znode_path)

        children_count = 0
        expired_count = 0
        removed_count = 0

        logger.info("Checking %s znode (%d children)...", znode_path, len(children))

        for child in children:
            child_path = os.path.join(znode_path, child)
            child_data = self.zk_client.get(child_path)[1]

            # If a child is also a parent - go deeper, but don't remove the parent itself
            if int(child_data.numChildren) > 0:
                result = self.remove_expired_children(child_path, dry_run)
                children_count += result['all']
                expired_count += result['expired']
                removed_count += result['removed']
                continue

            children_count += 1

            if self.is_znode_expired(child_data):
                expired_count += 1

                if dry_run:
                    logger.info('%s: (%s)', child_path, datetime.fromtimestamp(child_data.mtime // 1000))
                else:
                    removed_count += self.remove_znode(child_path)

        data = {'all': children_count, 'expired': expired_count, 'removed': removed_count}
        logger.info('Finished checking %s: %s' % (znode_path, data))

        return {
            'all': children_count,
            'expired': expired_count,
            'removed': removed_count
        }

    def remove_znode(self, znode_path):
        """ Removes znode and returns True in case of success. False otherwise. """
        try:
            self.zk_client.delete(znode_path)
            logger.info("Successfully removed %s node.", znode_path)
            return True
        except Exception as e:
            logger.error("Removing '%s' failed. Message: '%s'", znode_path, e)
            return False

    def tokens_cleanup(self, dry_run=False):
        all_znodes, expired_znodes, removed_znodes = 0, 0, 0
        self.zk_client.start(timeout=5)

        for parent_node in self.znodes_to_cleanup:
            result = self.remove_expired_children(parent_node, dry_run)
            all_znodes += result['all']
            expired_znodes += result['expired']
            removed_znodes += result['removed']

        logger.info("Cleanup finished. Found %d expired znodes out of %d. Removed %d.", expired_znodes, all_znodes, removed_znodes)

        self.zk_client.stop()

def main():
    """
    Usage:
        python zookeeper-cleanup.py -c <vault_config_path> -H <zookeeper_host_and_port> [-d]
    """
    # Parse arguments
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', action='store', dest='vault_config', default=VAULT_DEFAULT_CONFIG_FILE_PATH,
                        help='Vault configuration file path')
    parser.add_argument('-d', action='store_true', dest='dry_run', default=False,
                        help='Whether to remove objects or just list them.')
    parser.add_argument('-H', action='store', dest='zookeeper_hosts',
                        help='Zookeeper <host1>:<port1>,<host2>:<port2> comma-separated list.', default=VAULT_DEFAULT_HOST)
    parser.add_argument('-t', action='store', dest='ttl', type=int, default=None,
                        help="Expiration TTL in seconds (if not specified or equal to 0, it is taken from vault's config).")
    parser.add_argument('-z', action='store', dest='paths_to_cleanup', default=ZOOKEEPER_DEFAULT_CLEANUP_PATHS,
                        help="Comma-separated paths that will be checked recursively.")
    parser.add_argument('--version', action='version', version='%(prog)s 1.0')
    args = parser.parse_args()

    # Get effective TTL (either from args or from config file)
    try:
        vault_ttl = args.ttl or get_ttl_from_vault_config(args.vault_config)
    except IOError:
        logger.error(
            "Vault config file (%s) doesn't exist. Provide a correct path or a TTL time (when -t is passed, Vault config isn't checked).",
            args.vault_config
        )
        exit(1)

    logger.info("Start Zookeeper Vault Tokens cleanup procedure (TTL = %ds). DRY_RUN=%s", vault_ttl, args.dry_run)
    cleanup_manager = ZookeeperCleanupManager(args.zookeeper_hosts, args.paths_to_cleanup, vault_ttl)
    cleanup_manager.tokens_cleanup(args.dry_run)


if __name__ == "__main__":
    main()

#!/bin/bash

# Usage:
#   $> regenerate_root_token.sh [<installation_dir>] [https://<vault_host>:<vault_port>]

INSTALLATION_DIR=${1:-/opt/workfusion}
VAULT_KEYS_FILE=$INSTALLATION_DIR/vault/keys/vault_keys.json
VAULT_HOST=${2:-https://127.0.0.1:8200}
HEADERS=(--header "Content-Type: application/json")

# Delete an existing token generation attempt
curl -k -X DELETE "$VAULT_HOST/v1/sys/generate-root/attempt"

# Initializing token generation
result=$(curl -k -s -X PUT "${HEADERS[@]}" "$VAULT_HOST/v1/sys/generate-root/attempt")
nonce=$(echo "$result" | jq .nonce | tr -d '"')
otp=$(echo "$result" | jq .otp | tr -d '"')
echo " => Generation of the token initialized successfully."

# Updating unseal keys to the procedure
keys=$(cat "$VAULT_KEYS_FILE" | jq -c .keys[] | tr -d '"' | tr '\n' ' ')
for key in $keys; do
    result=$(curl -k -s -X PUT "${HEADERS[@]}" -d "{\"key\":\"$key\",\"nonce\":\"$nonce\"}" "$VAULT_HOST/v1/sys/generate-root/update")
done

# Decoding root token. XOR between b64-decoded encoded_root_token and otp
encoded_root_token=$(echo "$result" | jq .encoded_root_token | tr -d '"')
root_token=$(python -c "from base64 import b64decode as decode; \
    print(bytearray([_a ^ _b for _a, _b in zip(bytearray(decode('${encoded_root_token}==')), bytearray('${otp}'.encode()))]).decode());")
echo " => Successfully generated root token."

# Backup old VAULT_KEYS_FILE
vault_keys_file_backup=$VAULT_KEYS_FILE.backup.$(date +'%Y-%m-%d_%H%M%S')
cp "$VAULT_KEYS_FILE" "$vault_keys_file_backup"
echo " => Backuped $VAULT_KEYS_FILE to $vault_keys_file_backup."

# Update VAULT_KEYS_FILE with the new root token
current_vault_keys=$(cat "$VAULT_KEYS_FILE")
jq '.root_token = $newToken' --arg newToken "$root_token" -c <<<"$current_vault_keys" > "$VAULT_KEYS_FILE"
echo " => Root token ($root_token) has been written to $VAULT_KEYS_FILE. Finishing."

6. Add execution permissions to the regenerate_root_token Bash script.

   $ chmod +x /tmp/regenerate_root_token.sh
   

7. Install the curl and jq packages if you don't have them already in the system.

Clean up expired tokens

The zookeeper-cleanup.py script uses the Python that comes with the WorkFusion package. To use it, source the environment file with the source /opt/workfusion/environment.sh command before triggering it.

$ source /opt/workfusion/environment.sh

You can only display expired tokens and list them with last-modification-datetime without deleting them. This process is called DRY_RUN, and you can run it first to check that everything is working (Python libraries, connection, and so on):

$ python /tmp/zookeeper-cleanup.py -d

The first line of the script's output shows the current TTL (from the Vault configuration unless provided as the -t argument) and whether it's a DRY_RUN or not. There are default paths to check for expired tokens (/vault/sys/token/id and /vault/sys/expire/id/auth/cert/login), but you can change them with the -z argument:

$ python /tmp/zookeeper-cleanup.py -d -z /some/another/path

If everything works fine, remove the expired tokens by executing any of the following commands:

$ python /tmp/zookeeper-cleanup.py                              # Everything default (localhost:2181, TTL from Vault's config)
$ python /tmp/zookeeper-cleanup.py -t 120                       # Default host, 120 seconds of expiration threshold
$ python /tmp/zookeeper-cleanup.py -c /opt/vault/config.hcl     # Everything default, TTL read from different Vault's config

After this operation, all expired tokens are removed. To verify this, rerun the script in the DRY_RUN mode:

$ python /tmp/zookeeper-cleanup.py -d

Re-generate Vault root token

The cleanup procedure is harmless for the authentication between WorkFusion's components and Vault. But it breaks the Root Token authentication used for initializing, installing, and configuring Vault via Ansible Installer (./install.sh <command> vault commands).

To fix this, do the following:

1. Re-generate the root token and add it to a file used by the Ansible installer.

   $ /tmp/regenerate_root_token.sh                                             # All defaults
   $ /tmp/regenerate_root_token.sh /different/path/to/vault_keys_file.json     # Custom file path with keys (used by the Ansible installer)
   

   The new token is written to the JSON file on the INT server. The default path is /opt/workfusion/vault/keys/vault_keys.json.

2. Open /opt/workfusion/vault/keys/vault_keys.json and copy the file's content. You will need these details to insert into ZooKeeper later.

3. Once new tokens are created in the vault_key.json file, update them tokens in ZooKeeper by running the commands:

   $ cd /opt/workfusion/zookeeper/bin
   $ ./zkCli.sh
   get /vault-init-keys
   

   As an output, you will see the key and token parts.

4. Run the following command to update the token:

   get /vault-init-keys '{"keys":["key1","key2","key3","key4","key5"],"keys_base64":["key base64 formate"],"root_token":"tokenvault"}'