#!/usr/bin/python3
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright 2017,2024 Lenovo
#
# 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.


import asyncio
import getpass
import optparse
import sys
import os
path = os.path.dirname(os.path.realpath(__file__))
path = os.path.realpath(os.path.join(path, '..', 'lib', 'python'))
if path.startswith('/opt'):
    # if installed into system path, do not muck with things
    sys.path.append(path)
import confluent.config.configmanager as cfm
import confluent.config.conf as conf
import confluent.main as main
import confluent.noderange as noderange

argparser = optparse.OptionParser(
    usage="Usage: %prog [options] [dump|restore|merge] [path]\n"
          "       %prog [-u] [-v] showattrib <noderange> <attribute>...")
argparser.add_option('-p', '--password',
                     help='Password to use to protect/unlock a protected dump')
argparser.add_option('-i', '--interactivepassword', help='Prompt for password',
                     action='store_true')
argparser.add_option('-r', '--redact', action='store_true',
                     help='Redact potentially sensitive data rather than store')
argparser.add_option('-u', '--unprotected', action='store_true',
                      help='Specify that no password should be used to protect'
                           ' the key information.  Fields will be encrypted, '
                           'but keys.json will contain unencrypted decryption'
                           ' keys that may be used to read the dump.  With '
                           'showattrib, show secret.* values decrypted to '
                           'plaintext and crypted.* values as their stored '
                           'hashes rather than masking them')
argparser.add_option('-s', '--skipkeys', action='store_true',
                      help='This specifies to dump the encrypted data without '
                           'dumping the keys needed to decrypt it.  This is '
                           'suitable for an automated incremental backup, '
                           'where an earlier password protected dump has a '
                           'protected keys.json file, and only the protected '
                           'data is needed.  keys do not change and as such '
                           'they do not require incremental backup')
argparser.add_option('-y', '--yaml', action='store_true',
                      help='Use YAML instead of JSON as file format')
argparser.add_option('-v', '--value-only', dest='valueonly',
                     action='store_true',
                     help='With showattrib, print only the values '
                          'without node and attribute names')
(options, args) = argparser.parse_args()


def _dispval(value):
    # Render a single scalar for display: decode bytes (e.g. decrypted
    # secrets) and stringify anything else, including non-string values
    if isinstance(value, bytes):
        return value.decode('utf-8', errors='replace')
    return str(value)


def show_attrib(options, args):
    if len(args) < 3:
        sys.stderr.write('Usage: confluentdbutil showattrib [-u] [-v] '
                         '<noderange> <attribute> [<attribute>...]\n')
        sys.exit(1)
    nr = args[1]
    requested = args[2:]
    # Load the config store and unlock the master key, mirroring the daemon
    # startup sequence in confluent.main.  _initsecurity handles the
    # externalcfgkey case (key protected by the contents of a key file);
    # otherwise the key is stored unprotected and init_masterkey unlocks it
    # without any password.  autogen=False keeps this read-only (never create a
    # key).
    cfm.init()
    try:
        main._initsecurity(conf.get_config())
        if not cfm._masterkey:
            cfm.init_masterkey(autogen=False)
    except Exception:
        sys.stderr.write("Error unlocking credential store\n")
        sys.exit(1)
    if not cfm._masterkey:
        sys.stderr.write("Error unlocking credential store\n")
        sys.exit(1)
    cfg = cfm.ConfigManager(None)
    try:
        nodes = noderange.NodeRange(nr, config=cfg).nodes
    except Exception as e:
        sys.stderr.write("Invalid Noderange: " + str(e) + "\n")
        sys.exit(1)
    attrs = cfg.get_node_attributes(list(nodes), requested,
                                    decrypt=options.unprotected)
    # Order nodes and attributes exactly as the server/nodeattrib do:
    # natural sort via humanify_nodename (identical to the client's
    # naturalize_string), with the same TypeError fallback as core.py.
    nodelist = list(attrs)
    try:
        nodelist.sort(key=noderange.humanify_nodename)
    except TypeError:
        nodelist.sort()
    for node in nodelist:
        nodeattrs = attrs[node]
        anames = list(nodeattrs)
        try:
            anames.sort(key=noderange.humanify_nodename)
        except TypeError:
            anames.sort()
        for aname in anames:
            adata = nodeattrs[aname]
            # Format like nodeattrib: a plain, list, or (decrypted) secret.*
            # value lives in 'value'; some attributes come back as a bare list;
            # crypted.* is a one-way hash stored as 'hashvalue'.  Without -u,
            # sensitive values are masked the same way nodeattrib does.
            if isinstance(adata, (list, tuple)):
                val = ','.join(_dispval(x) for x in adata)
            elif 'value' in adata:
                aval = adata['value']
                if aval is None:
                    val = ''
                elif isinstance(aval, (list, tuple)):
                    val = ','.join(_dispval(x) for x in aval)
                else:
                    val = _dispval(aval)
            elif 'hashvalue' in adata:
                if options.unprotected:
                    val = _dispval(adata['hashvalue'])
                else:
                    val = '******** (hidden: use -u to show)'
            elif 'cryptvalue' in adata:
                val = '******** (hidden: use -u to show)'
            else:
                val = ''
            if options.valueonly:
                print(val)
            else:
                print('{0}: {1}: {2}'.format(node, aname, val))
        # Report explicitly named (non-glob) attributes that are not set.
        for attr in requested:
            if '*' not in attr and attr not in nodeattrs:
                if options.valueonly:
                    print('')
                else:
                    print('{0}: {1}:'.format(node, attr))
    sys.exit(0)


if args and args[0] == 'showattrib':
    show_attrib(options, args)
if len(args) != 2 or args[0] not in ('dump', 'restore', 'merge'):
    argparser.print_help()
    sys.exit(1)
dumpdir = args[1]


if args[0] in ('restore', 'merge'):
    pid = main.is_running()
    if pid is not None:
        print("Confluent is running, must shut down to restore db")
        sys.exit(1)
    stinf = os.stat('/etc/confluent')
    owner = stinf.st_uid
    group = stinf.st_gid
    password = options.password
    if options.interactivepassword:
        password = getpass.getpass('Enter password to restore backup: ')
    try:
        stateless = args[0] == 'restore'
        cfm.init(stateless)
        cfm.statelessmode = stateless
        skipped = {'nodes': [], 'nodegroups': []}
        
        # Use the format parameter based on the --yaml option
        format = 'yaml' if options.yaml else 'json'
        
        dp = cfm.restore_db_from_directory(
            dumpdir, password,
            merge="skip" if args[0] == 'merge' else False, 
            skipped=skipped,
            format=format)
            
        asyncio.run(dp)
        if skipped['nodes']:
            skippedn = ','.join(skipped['nodes'])
            print('The following nodes were skipped during merge: '
                  '{}'.format(skippedn))
        if skipped['nodegroups']:
            skippedn = ','.join(skipped['nodegroups'])
            print('The following node groups were skipped during merge: '
                  '{}'.format(skippedn))
        cfm.statelessmode = False
        cfm.ConfigManager.wait_for_sync(True)
        if owner != 0:
            for targdir in os.walk('/etc/confluent'):
                os.chown(targdir[0], owner, group)
                for f in targdir[2]:
                    os.chown(os.path.join(targdir[0], f), owner, group)
    except Exception as e:
        print(str(e))
        sys.exit(1)
elif args[0] == 'dump':
    password = options.password
    if not password and options.interactivepassword:
        passcfm = None
        while passcfm is None or password != passcfm:
            password = getpass.getpass(
                'Enter password to protect the backup: ')
            passcfm = getpass.getpass('Confirm password to protect the backup: ')
    if password is None and not (options.unprotected or options.redact
                                         or options.skipkeys):
        print("Must indicate a password to protect or -u to opt opt of "
              "secure value protection or -r to redact sensitive information, "
              "or -s to do encrypted backup that requires keys.json from "
              "another backup to restore.")
        sys.exit(1)
    os.umask(0o77)
    main._initsecurity(conf.get_config())
    if not os.path.exists(dumpdir):
        os.makedirs(dumpdir)
    
    # Use the format parameter based on the --yaml option
    format = 'yaml' if options.yaml else 'json'
    dp = cfm.dump_db_to_directory(dumpdir, password, options.redact,
                                  options.skipkeys, format=format)
    asyncio.run(dp)


