#!/usr/bin/python3.11
# -*- coding: utf-8 -*-
########### section II below by NTESS ################################
license_ntess="""
Copyright (c) 2010-2018 National Technology & Engineering Solutions
of Sandia, LLC (NTESS). Under the terms of Contract DE-NA0003525 with
NTESS, the U.S. Government retains certain rights in this software.
Copyright (c) 2010-2018 Open Grid Computing, Inc. All rights reserved.

This software is available to you under a choice of one of two
licenses.  You may choose to be licensed under the terms of the GNU
General Public License (GPL) Version 2, available from the file
COPYING in the main directory of this source tree, or the BSD-type
license below:

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:

     Redistributions of source code must retain the above copyright
     notice, this list of conditions and the following disclaimer.

     Redistributions in binary form must reproduce the above
     copyright notice, this list of conditions and the following
     disclaimer in the documentation and/or other materials provided
     with the distribution.

     Neither the name of Sandia nor the names of any contributors may
     be used to endorse or promote products derived from this software
     without specific prior written permission.

     Neither the name of Open Grid Computing nor the names of any
     contributors may be used to endorse or promote products derived
     from this software without specific prior written permission.

     Modified source versions must be plainly marked as such, and
     must not be misrepresented as being the original software.


THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

#  This script handles anonymization of csv data.

########### section I from published hostlist.py #####################
#
# Hostlist library
#
# Copyright (C) 2008-2018
#                    Kent Engström <kent@nsc.liu.se>,
#                    Thomas Bellman <bellman@nsc.liu.se>,
#                    Pär Lindfors <paran@nsc.liu.se> and
#                    Torbjörn Lönnemark <ketl@nsc.liu.se>,
#                    National Supercomputer Centre
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.

"""Handle hostlist expressions.

This module provides operations to expand and collect hostlist
expressions.

The hostlist expression syntax is the same as in several programs
developed at LLNL (https://computing.llnl.gov/linux/). However in
corner cases the behaviour of this module have not been compared for
compatibility with pdsh/dshbak/SLURM et al.
"""

__version__ = "1.18"

import re
import itertools

# Replace range with xrange on Python 2, do nothing on Python 3 (where xrange
# does not exist, and range returns an iterator)
try:
    range = xrange
except:
    pass

# Exception used for error reporting to the caller
class BadHostlist(Exception): pass

# Configuration to guard against ridiculously long expanded lists
MAX_SIZE = 100000

# Hostlist expansion

def expand_hostlist(hostlist, allow_duplicates=False, sort=False):
    """Expand a hostlist expression string to a Python list.

    Example: expand_hostlist("n[9-11],d[01-02]") ==>
             ['n9', 'n10', 'n11', 'd01', 'd02']

    Unless allow_duplicates is true, duplicates will be purged
    from the results. If sort is true, the output will be sorted.
    """

    results = []
    bracket_level = 0
    part = ""

    for c in hostlist + ",":
        if c == "," and bracket_level == 0:
            # Comma at top level, split!
            if part: results.extend(expand_part(part))
            part = ""
            bad_part = False
        else:
            part += c

        if c == "[": bracket_level += 1
        elif c == "]": bracket_level -= 1

        if bracket_level > 1:
            raise BadHostlist("nested brackets")
        elif bracket_level < 0:
            raise BadHostlist("unbalanced brackets")

    if bracket_level > 0:
        raise BadHostlist("unbalanced brackets")

    if not allow_duplicates:
        results = remove_duplicates(results)
    if sort:
        results = numerically_sorted(results)
    return results

def expand_part(s):
    """Expand a part (e.g. "x[1-2]y[1-3][1-3]") (no outer level commas)."""

    # Base case: the empty part expand to the singleton list of ""
    if s == "":
        return [""]

    # Split into:
    # 1) prefix string (may be empty)
    # 2) rangelist in brackets (may be missing)
    # 3) the rest

    m = re.match(r'([^,\[]*)(\[[^\]]*\])?(.*)', s)
    (prefix, rangelist, rest) = m.group(1,2,3)

    # Expand the rest first (here is where we recurse!)
    rest_expanded = expand_part(rest)

    # Expand our own part
    if not rangelist:
        # If there is no rangelist, our own contribution is the prefix only
        us_expanded = [prefix]
    else:
        # Otherwise expand the rangelist (adding the prefix before)
        us_expanded = expand_rangelist(prefix, rangelist[1:-1])

    # Combine our list with the list from the expansion of the rest
    # (but guard against too large results first)
    if len(us_expanded) * len(rest_expanded) > MAX_SIZE:
        raise BadHostlist("results too large")

    return [us_part + rest_part
            for us_part in us_expanded
            for rest_part in rest_expanded]

def expand_rangelist(prefix, rangelist):
    """ Expand a rangelist (e.g. "1-10,14"), putting a prefix before."""

    # Split at commas and expand each range separately
    results = []
    for range_ in rangelist.split(","):
        results.extend(expand_range(prefix, range_))
    return results

def expand_range(prefix, range_):
    """ Expand a range (e.g. 1-10 or 14), putting a prefix before."""

    # Check for a single number first
    m = re.match(r'^[0-9]+$', range_)
    if m:
        return ["%s%s" % (prefix, range_)]

    # Otherwise split low-high
    m = re.match(r'^([0-9]+)-([0-9]+)$', range_)
    if not m:
        raise BadHostlist("bad range")

    (s_low, s_high) = m.group(1,2)
    low = int(s_low)
    high = int(s_high)
    width = len(s_low)

    if high < low:
        raise BadHostlist("start > stop")
    elif high - low > MAX_SIZE:
        raise BadHostlist("range too large")

    results = []
    for i in range(low, high+1):
        results.append("%s%0*d" % (prefix, width, i))
    return results

def remove_duplicates(l):
    """Remove duplicates from a list (but keep the order)."""
    seen = set()
    results = []
    for e in l:
        if e not in seen:
            results.append(e)
            seen.add(e)
    return results

# Hostlist collection

def collect_hostlist(hosts, silently_discard_bad = False):
    """Collect a hostlist string from a Python list of hosts.

    We start grouping from the rightmost numerical part.
    Duplicates are removed.

    A bad hostname raises an exception (unless silently_discard_bad
    is true causing the bad hostname to be silently discarded instead).
    """

    # Split hostlist into a list of (host, "") for the iterative part.
    # (Also check for bad node names now)
    # The idea is to move already collected numerical parts from the
    # left side (seen by each loop) to the right side (just copied).

    left_right = []
    for host in hosts:
        # We remove leading and trailing whitespace first, and skip empty lines
        host = host.strip()
        if host == "": continue

        # We cannot accept a host containing any of the three special
        # characters in the hostlist syntax (comma and flat brackets)
        if re.search(r'[][,]', host):
            if silently_discard_bad:
                continue
            else:
                raise BadHostlist("forbidden character")

        left_right.append((host, ""))

    # Call the iterative function until it says it's done
    looping = True
    while looping:
        left_right, looping = collect_hostlist_1(left_right)
    return ",".join([left + right for left, right in left_right])

def collect_hostlist_1(left_right):
    """Collect a hostlist string from a list of hosts (left+right).

    The input is a list of tuples (left, right). The left part
    is analyzed, while the right part is just passed along
    (it can contain already collected range expressions).
    """

    # Scan the list of hosts (left+right) and build two things:
    # *) a set of all hosts seen (used later)
    # *) a list where each host entry is preprocessed for correct sorting

    sortlist = []
    remaining = set()
    for left, right in left_right:
        host = left + right
        remaining.add(host)

        # Match the left part into parts
        m = re.match(r'^(.*?)([0-9]+)?([^0-9]*)$', left)
        (prefix, num_str, suffix) = m.group(1,2,3)

        # Add the right part unprocessed to the suffix.
        # This ensures than an already computed range expression
        # in the right part is not analyzed again.
        suffix = suffix + right

        if num_str is None:
            # A left part with no numeric part at all gets special treatment!
            # The regexp matches with the whole string as the suffix,
            # with nothing in the prefix or numeric parts.
            # We do not want that, so we move it to the prefix and put
            # None as a special marker where the suffix should be.
            assert prefix == ""
            sortlist.append(((host, None), None, None, host))
        else:
            # A left part with at least an numeric part
            # (we care about the rightmost numeric part)
            num_int = int(num_str)
            num_width = len(num_str) # This width includes leading zeroes
            sortlist.append(((prefix, suffix), num_int, num_width, host))

    # Sort lexicographically, first on prefix, then on suffix, then on
    # num_int (numerically), then...
    # This determines the order of the final result.

    sortlist.sort()

    # We are ready to collect the result parts as a list of new (left,
    # right) tuples.

    results = []
    needs_another_loop = False

    # Now group entries with the same prefix+suffix combination (the
    # key is the first element in the sortlist) to loop over them and
    # then to loop over the list of hosts sharing the same
    # prefix+suffix combination.

    for ((prefix, suffix), group) in itertools.groupby(sortlist,
                                                       key=lambda x:x[0]):

        if suffix is None:
            # Special case: a host with no numeric part
            results.append(("", prefix)) # Move everything to the right part
            remaining.remove(prefix)
        else:
            # The general case. We prepare to collect a list of
            # ranges expressed as (low, high, width) for later
            # formatting.
            range_list = []

            for ((prefix2, suffix2), num_int, num_width, host) in group:
                if host not in remaining:
                    # Below, we will loop internally to enumate a whole range
                    # at a time. We then remove the covered hosts from the set.
                    # Therefore, skip the host here if it is gone from the set.
                    continue
                assert num_int is not None

                # Scan for a range starting at the current host
                low = num_int
                while True:
                    host = "%s%0*d%s" % (prefix, num_width, num_int, suffix)
                    if host in remaining:
                        remaining.remove(host)
                        num_int += 1
                    else:
                        break
                high = num_int - 1
                assert high >= low
                range_list.append((low, high, num_width))

            # We have a list of ranges to format. We make sure
            # we move our handled numerical part to the right to
            # stop it from being processed again.
            needs_another_loop = True
            if len(range_list) == 1 and range_list[0][0] == range_list[0][1]:
                # Special case to make sure that n1 is not shown as n[1] etc
                results.append((prefix,
                                "%0*d%s" %
                               (range_list[0][2], range_list[0][0], suffix)))
            else:
                # General case where high > low
                results.append((prefix, "[" + \
                                   ",".join([format_range(l, h, w)
                                             for l, h, w in range_list]) + \
                                   "]" + suffix))

    # At this point, the set of remaining hosts should be empty and we
    # are ready to return the result, together with the flag that says
    # if we need to loop again (we do if we have added something to a
    # left part).
    assert not remaining
    return results, needs_another_loop

def format_range(low, high, width):
    """Format a range from low to high inclusively, with a certain width."""

    if low == high:
        return "%0*d" % (width, low)
    else:
        return "%0*d-%0*d" % (width, low, width, high)

# Sort a list of hosts numerically

def numerically_sorted(l):
    """Sort a list of hosts numerically.

    E.g. sorted order should be n1, n2, n10; not n1, n10, n2.
    """

    return sorted(l, key=numeric_sort_key)

nsk_re = re.compile("([0-9]+)|([^0-9]+)")
def numeric_sort_key(x):
    return [handle_int_nonint(i_ni) for i_ni in nsk_re.findall(x)]

def handle_int_nonint(int_nonint_tuple):
    if int_nonint_tuple[0]:
        return int(int_nonint_tuple[0])
    else:
        return int_nonint_tuple[1]

# Parse SLURM_TASKS_PER_NODE into a list of task numbers
#
# Description from the SLURM sbatch man page:
#              Number of tasks to be initiated on each node. Values
#              are comma separated and in the same order as
#              SLURM_NODELIST.  If two or more consecutive nodes are
#              to have the same task count, that count is followed by
#              "(x#)" where "#" is the repetition count. For example,
#              "SLURM_TASKS_PER_NODE=2(x3),1" indicates that the first
#              three nodes will each execute three tasks and the
#              fourth node will execute one task.

def parse_slurm_tasks_per_node(s):
    res = []
    for part in s.split(","):
        m = re.match(r'^([0-9]+)(\(x([0-9]+)\))?$', part)
        if m:
            tasks = int(m.group(1))
            repetitions = m.group(3)
            if repetitions is None:
                repetitions = 1
            else:
                repetitions = int(repetitions)
            if repetitions > MAX_SIZE:
                raise BadHostlist("task list repetitions too large")
            for i in range(repetitions):
                res.append(tasks)
        else:
            raise BadHostlist("bad task list syntax")
    return res

########### end of section I from hostlist.py #####################

########### section II from NTESS #####################
#
# dump anonymized csv files and the substitution maps.
# substitutions are globally consistent for the batch of files.
# substitution maps may be reused and extended across batches.
#
#
import argparse
import os.path
import sys
import random

ioffset = 1000000001
def map_int(val, vmap):
    """return mapping of int to range 0-N of N unique values seen"""
    if not val in vmap:
        x = str(len(vmap)+ioffset)
        vmap[val] = x
        return x
    else:
        return vmap[val]

def map_path(val, vmap):
    """replace value in elts[col] with anonymized one. vmap is pairings"""
    x = val.split("/")
    y = []
    for i in x:
        if len(i) > 0:
            if not i in vmap:
                j = "p" + str(len(vmap)+1)
                vmap[i] = j
                y.append(j)
            else:
                y.append(vmap[i])
        else:
            y.append(i)
    return "/".join(y)

def map_name(val, vmap):
    """replace value in elts[col] with anonymized one. vmap is pairings"""
    if not val in vmap:
        x = "n" + str(len(vmap)+1)
        vmap[val] = x
        return x
    else:
        return vmap[val]

def map_host(val, vmap):
    """replace value(s) in elts[col] with anonymized one. vmap is pairings.
    This is specific to hostname and host lists with manual pre-mapping.
    hostnames look like: prefixNNN, prefixNNN-netsuffix where NNN may be replaced
    with comma/- separated range of numbers in [].
    special pairs are:
    netdomains <subs> where subs might be "ca.sandia.gov,sandia.gov"
    netdomains must be in order of decreasing detail.
    """
    if val in vmap:
        return vmap[val]
    # get rid of range notation
    x = expand_hostlist(val)
    y = []
    #fix individual names
    for i in x:
        p = host_substitute(i, vmap)
        y.append(p)
    z = collect_hostlist(y)
    return z

def host_substitute(h, vmap):
    """ replace - or number trailed parts of h. Something must change or error,
    unless name starts with localhost."""
    if h.startswith("localhost"):
        vmap[h] = h
        return h
    # delete fqn suffixes
    doms = vmap["netdomains"]
    for k in doms:
        if k in h:
            h = h.replace(k,"")
    # replace name parts
    cparts = h.split("-")
    canon = []
    for i in cparts:
        # preserve empty list elements
        if not i:
            canon.append(i)
            continue
        # direct replacement
        if i in vmap:
            canon.append(vmap[i])
            continue
        # number suffixed replacement
        nonum = i.rstrip("0123456789")
        if len(nonum) != len(i):
            if nonum in vmap:
                head = vmap[nonum]
                head += i[len(nonum):]
                vmap[i] = head
                canon.append(head)
                continue
            else:
                print(f"Fix hmap. Can't map {nonum} in {i} of {h}")
                sys.exit(1)
        else:
            print(f"Fix hmap file. Can't map {i} in {h}")
            sys.exit(1)

    fin = "-".join(canon)
    if not h in vmap:
        vmap[h] = fin
    return fin

def map_header(val, meth):
    return "anonymized_" + meth + "_" + val

mapmagic = "#anonymize-csv-map"

def write_map(maps, kind, fn):
    if len(maps[kind]) > 0:
        with open(fn, "w") as o:
            print(mapmagic, kind, file=o)
            m = maps[kind]
            for key in sorted(m):
                if key == "netdomains":
                    if len(m[key]) > 0:
                        print(key, ",".join(m[key]), file=o)
                else:
                    print(key, m[key], file=o)

def reload_map(m, fn, kind):
    """prime m with data from prior run. input can be empty, otherwise must match kind."""
    header = True
    with open(fn, "r") as i:
        for ln in i:
            ln = ln.rstrip()
            if len(ln) < 1:
                continue
            (k,v) = ln.split()
            if header:
                if k != mapmagic:
                    print(f"Error loading map- not a mapping file: {fn}")
                    sys.exit(1)
                if v != kind:
                    print(f"Error loading map file- wrong type. {kind} vs {v} in")
                    sys.exit(1)
                header = False
                continue
            m[k] = v

def convert_col(kstr):
    """convert extended cut index string to python notation."""
    k = int(kstr)
    if k == 0:
        print("column numbers are 1-N (as cut(1)) or negative, not 0.")
        sys.exit(1)
    if k > 0:
        k -= 1
    return k

rewrite = {
        "int" : map_int,
        "path" : map_path,
        "name" : map_name,
        "host" : map_host,
    }

def gen_args(colfile, sep):
    cf = colfile.split(",")
    if len(cf) < 2:
       sys.stderr.write("gen-args needs (kind:metricname)*,filename\n")
       sys.exit(1)
    metrics = cf[0:-1]
    fname = cf[-1]
    with open(fname, "r") as i:
        for ln in i:
            if ln[0] != '#':
                sys.stderr.write("cannot parse header line in " +  fname + "\n")
                sys.exit(1)
            ln = ln[1:].rstrip()
            cols = ln.split(sep)
            s = []
            for j in metrics:
                (k, m) = j.split(":")
                if m in cols:
                    c = cols.index(m) + 1
                    s.append(":".join([k,str(c)]))
                else:
                    sys.stderr.write("metric " + m + " not in " +  fname + "\n")
                    sys.stderr.write( str(cols) + "\n")
                    sys.exit(1)
            for j in s:
                sys.stdout.write(j+" ")
            sys.stdout.write("\n")
            break

if __name__ == "__main__":
    """Anonymize columns, e.g.
    ./anonymize-csv  "--c=|" --se 123 2:int 3:name 5:path file*
    #time,wheels,model,year,dblocation
    1550000000,3,tricycle,1972,/var/cardb
    1550000000,3,tricycle,1972,/var/cardc
    -->
    #time,wheels,model,year,dblocation
    1550000000,1000000001,n1,1972,/p1/p2
    1550000000,1000000001,n1,1972,/p1/p3
    """
    parser = argparse.ArgumentParser(description="Anonymize columnar data files.")
    parser.add_argument('mappings', metavar='M:C', nargs='*',
            help='a mapping:column number pair or filename. M is one of int,path,name,host. C is a nonzero number. Negative number count back from last column.')
    parser.add_argument("--input", default=None, help="input files", action="append", metavar="csv-file")
    parser.add_argument("--out-dir", default=None, help="output files location")
    parser.add_argument("--col-sep", default=None, help="column separator")
    parser.add_argument("--seed", default=None, type=int, help="randomization seed")
    parser.add_argument("--save-maps", default=None, help="prefix for map output files")
    parser.add_argument("--imap", default=None, help="input file of initial integer mapping.")
    parser.add_argument("--nmap", default=None, help="input file of initial name mapping.")
    parser.add_argument("--pmap", default=None, help="input file of initial path element mapping.")
    parser.add_argument("--hmap", default=None, help="input file of complete hostname mapping.")
    parser.add_argument("--gen-args", default=None, help="input metric names and file")
    parser.add_argument("--debug", default=False, action='store_true', help="enable debug")
    if len(sys.argv) < 2:
        print("Need some arguments. Try -h")
        sys.exit(1)
    args = parser.parse_args()

    if args.debug:
        print(args)
    sep = args.col_sep
    if not sep:
        sep = ','
    if args.gen_args:
        gen_args(args.gen_args, sep)
        sys.exit(0)
    files = []
    for i in args.input:
        files.extend(i.split())
    cols = []
    if len(args.mappings) == 0:
        print("Nothing to do")
        sys.exit(1)
    for i in args.mappings:
        if ":" in i:
            cols.append(i)
        else:
            files.append(i)
    if args.debug:
        print(files)
        print(cols)
    maps = dict()
    for key in list(rewrite.keys()):
        maps[key] =  dict()
    if args.imap:
        reload_map(maps["int"], args.imap, "int")
    if args.nmap:
        reload_map(maps["name"], args.nmap, "name")
    if args.pmap:
        reload_map(maps["path"], args.pmap, "path")
    maps["host"]["netdomains"] = []
    if args.hmap:
        reload_map(maps["host"], args.hmap, "host")
        if "netdomains" in maps["host"]:
            maps["host"]["netdomains"] = maps["host"]["netdomains"].split(",")

    prefix = args.save_maps
    if not prefix:
        prefix = 'anonmap_'
    seed = args.seed
    if not seed:
        seed = random.randint(1,sys.maxsize)
    random.seed(seed)
    od = args.out_dir
    if not od or not os.path.isdir(od):
        print("output directory not given or not found")
        sys.exit(1)
    for f in files:
        if not os.path.isfile(f):
            print("File not found {f}")
            sys.exit(1)
    for f in files:
        if args.debug:
            print(f)
        bname = os.path.basename(f)
        with open(f, "r") as i:
            out = os.path.join(od, bname)
            with open(out, "w") as o:
                for ln in i:
                    ln = ln.rstrip()
                    if len(ln) < 3:
                        print(ln, file=o)
                        continue
                    if ln[0] == '#':
                        x = ln.split(sep)
                        for c in cols:
                            (method, k) = c.split(":")
                            k = convert_col(k)
                            x[k] = map_header(x[k], method)
                        y = sep.join(x)
                        print(y, file=o)
                        continue
                    x = ln.split(sep)
                    for c in cols:
                        (method, k) = c.split(":")
                        if method == "host" and args.hmap == None:
                            print("--hmap required for host mapping")
                            sys.exit(1)
                        k = convert_col(k)
                        x[k] = rewrite[method](x[k], maps[method])
                    y = sep.join(x)
                    print(y, file=o)


    # print seed file if we ever use random

    # dump imap, pmap, nmap, pmap
    write_map(maps, "int", os.path.join(od, prefix + "imap.txt"))
    write_map(maps, "path", os.path.join(od, prefix + "pmap.txt"))
    write_map(maps, "name", os.path.join(od, prefix + "nmap.txt"))
    write_map(maps, "host", os.path.join(od, prefix + "hmap.txt"))

