/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package io.opentelemetry.instrumentation.api.db;

import static java.util.Arrays.asList;
import static java.util.Collections.unmodifiableMap;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;

/**
 * This class is responsible for masking potentially sensitive data in Redis commands.
 *
 * <p>Examples:
 *
 * <table>
 *   <tr>
 *     <th>Raw command</th>
 *     <th>Normalized command</th>
 *   </tr>
 *   <tr>
 *     <td>{@code AUTH password}</td>
 *     <td>{@code AUTH ?}</td>
 *   </tr>
 *   <tr>
 *     <td>{@code HMSET hash creditcard 1234567887654321 address asdf}</td>
 *     <td>{@code HMSET hash creditcard ? address ?}</td>
 *   </tr>
 * </table>
 */
public final class RedisCommandSanitizer {

  private static final Map<String, CommandSanitizer> SANITIZERS;
  private static final CommandSanitizer DEFAULT = new CommandAndNumArgs(0);

  static {
    Map<String, CommandSanitizer> sanitizers = new HashMap<>();

    CommandSanitizer keepOneArg = new CommandAndNumArgs(1);
    CommandSanitizer keepTwoArgs = new CommandAndNumArgs(2);
    CommandSanitizer setMultiHashField = new MultiKeyValue(1);
    CommandSanitizer setMultiField = new MultiKeyValue(0);

    // Cluster
    for (String command : asList("CLUSTER", "FAILOVER", "READONLY", "READWRITE")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Connection
    sanitizers.put("AUTH", DEFAULT);
    // HELLO can contain AUTH data
    sanitizers.put("HELLO", keepTwoArgs);
    for (String command : asList("CLIENT", "ECHO", "PING", "QUIT", "RESET", "SELECT")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Geo
    for (String command :
        asList(
            "GEOADD",
            "GEODIST",
            "GEOHASH",
            "GEOPOS",
            "GEORADIUS",
            "GEORADIUS_RO",
            "GEORADIUSBYMEMBER",
            "GEORADIUSBYMEMBER_RO",
            "GEOSEARCH",
            "GEOSEARCHSTORE")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Hashes
    sanitizers.put("HMSET", setMultiHashField);
    sanitizers.put("HSET", setMultiHashField);
    sanitizers.put("HSETNX", keepTwoArgs);
    for (String command :
        asList(
            "HDEL",
            "HEXISTS",
            "HGET",
            "HGETALL",
            "HINCRBY",
            "HINCRBYFLOAT",
            "HKEYS",
            "HLEN",
            "HMGET",
            "HRANDFIELD",
            "HSCAN",
            "HSTRLEN",
            "HVALS")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // HyperLogLog
    sanitizers.put("PFADD", keepOneArg);
    for (String command : asList("PFCOUNT", "PFMERGE")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Keys
    // MIGRATE can contain AUTH data
    sanitizers.put("MIGRATE", new CommandAndNumArgs(6));
    sanitizers.put("RESTORE", keepTwoArgs);
    for (String command :
        asList(
            "COPY",
            "DEL",
            "DUMP",
            "EXISTS",
            "EXPIRE",
            "EXPIREAT",
            "EXPIRETIME",
            "KEYS",
            "MOVE",
            "OBJECT",
            "PERSIST",
            "PEXPIRE",
            "PEXPIREAT",
            "PEXPIRETIME",
            "PTTL",
            "RANDOMKEY",
            "RENAME",
            "RENAMENX",
            "SCAN",
            "SORT",
            "SORT_RO",
            "TOUCH",
            "TTL",
            "TYPE",
            "UNLINK",
            "WAIT")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Lists
    sanitizers.put("LINSERT", keepTwoArgs);
    sanitizers.put("LPOS", keepOneArg);
    sanitizers.put("LPUSH", keepOneArg);
    sanitizers.put("LPUSHX", keepOneArg);
    sanitizers.put("LREM", keepOneArg);
    sanitizers.put("LSET", keepOneArg);
    sanitizers.put("RPUSH", keepOneArg);
    sanitizers.put("RPUSHX", keepOneArg);
    for (String command :
        asList(
            "BLMOVE",
            "BLMPOP",
            "BLPOP",
            "BRPOP",
            "BRPOPLPUSH",
            "LINDEX",
            "LLEN",
            "LMOVE",
            "LMPOP",
            "LPOP",
            "LRANGE",
            "LTRIM",
            "RPOP",
            "RPOPLPUSH")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Pub/Sub
    sanitizers.put("PUBLISH", keepOneArg);
    for (String command :
        asList(
            "PSUBSCRIBE",
            "PUBSUB",
            "PUNSUBSCRIBE",
            "SPUBLISH",
            "SSUBSCRIBE",
            "SUBSCRIBE",
            "SUNSUBSCRIBE",
            "UNSUBSCRIBE")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Scripting
    sanitizers.put("EVAL", Eval.INSTANCE);
    sanitizers.put("EVAL_RO", Eval.INSTANCE);
    sanitizers.put("EVALSHA", Eval.INSTANCE);
    sanitizers.put("EVALSHA_RO", Eval.INSTANCE);
    sanitizers.put("SCRIPT", KeepAllArgs.INSTANCE);

    // Server
    // CONFIG SET can set any property, including the master password
    sanitizers.put("CONFIG", keepTwoArgs);
    for (String command :
        asList(
            "ACL",
            "BGREWRITEAOF",
            "BGSAVE",
            "COMMAND",
            "DBSIZE",
            "DEBUG",
            "FLUSHALL",
            "FLUSHDB",
            "INFO",
            "LASTSAVE",
            "LATENCY",
            "LOLWUT",
            "MEMORY",
            "MODULE",
            "MONITOR",
            "PSYNC",
            "REPLICAOF",
            "ROLE",
            "SAVE",
            "SHUTDOWN",
            "SLAVEOF",
            "SLOWLOG",
            "SWAPDB",
            "SYNC",
            "TIME")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Sets
    sanitizers.put("SADD", keepOneArg);
    sanitizers.put("SISMEMBER", keepOneArg);
    sanitizers.put("SMISMEMBER", keepOneArg);
    sanitizers.put("SMOVE", keepTwoArgs);
    sanitizers.put("SREM", keepOneArg);
    for (String command :
        asList(
            "SCARD",
            "SDIFF",
            "SDIFFSTORE",
            "SINTER",
            "SINTERCARD",
            "SINTERSTORE",
            "SMEMBERS",
            "SPOP",
            "SRANDMEMBER",
            "SSCAN",
            "SUNION",
            "SUNIONSTORE")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Sorted Sets
    sanitizers.put("ZADD", keepOneArg);
    sanitizers.put("ZCOUNT", keepOneArg);
    sanitizers.put("ZINCRBY", keepOneArg);
    sanitizers.put("ZLEXCOUNT", keepOneArg);
    sanitizers.put("ZMSCORE", keepOneArg);
    sanitizers.put("ZRANGEBYLEX", keepOneArg);
    sanitizers.put("ZRANGEBYSCORE", keepOneArg);
    sanitizers.put("ZRANK", keepOneArg);
    sanitizers.put("ZREM", keepOneArg);
    sanitizers.put("ZREMRANGEBYLEX", keepOneArg);
    sanitizers.put("ZREMRANGEBYSCORE", keepOneArg);
    sanitizers.put("ZREVRANGEBYLEX", keepOneArg);
    sanitizers.put("ZREVRANGEBYSCORE", keepOneArg);
    sanitizers.put("ZREVRANK", keepOneArg);
    sanitizers.put("ZSCORE", keepOneArg);
    for (String command :
        asList(
            "BZMPOP",
            "BZPOPMAX",
            "BZPOPMIN",
            "ZCARD",
            "ZDIFF",
            "ZDIFFSTORE",
            "ZINTER",
            "ZINTERCARD",
            "ZINTERSTORE",
            "ZMPOP",
            "ZPOPMAX",
            "ZPOPMIN",
            "ZRANDMEMBER",
            "ZRANGE",
            "ZRANGESTORE",
            "ZREMRANGEBYRANK",
            "ZREVRANGE",
            "ZSCAN",
            "ZUNION",
            "ZUNIONSTORE")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Streams
    sanitizers.put("XADD", new MultiKeyValue(2));
    for (String command :
        asList(
            "XACK",
            "XAUTOCLAIM",
            "XCLAIM",
            "XDEL",
            "XGROUP",
            "XINFO",
            "XLEN",
            "XPENDING",
            "XRANGE",
            "XREAD",
            "XREADGROUP",
            "XREVRANGE",
            "XTRIM")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Strings
    sanitizers.put("APPEND", keepOneArg);
    sanitizers.put("GETSET", keepOneArg);
    sanitizers.put("MSET", setMultiField);
    sanitizers.put("MSETNX", setMultiField);
    sanitizers.put("PSETEX", keepTwoArgs);
    sanitizers.put("SET", keepOneArg);
    sanitizers.put("SETEX", keepTwoArgs);
    sanitizers.put("SETNX", keepOneArg);
    sanitizers.put("SETRANGE", keepOneArg);
    for (String command :
        asList(
            "BITCOUNT",
            "BITFIELD",
            "BITFIELD_RO",
            "BITOP",
            "BITPOS",
            "DECR",
            "DECRBY",
            "GET",
            "GETBIT",
            "GETDEL",
            "GETEX",
            "GETRANGE",
            "INCR",
            "INCRBY",
            "INCRBYFLOAT",
            "LCS",
            "MGET",
            "SETBIT",
            "STRALGO",
            "STRLEN")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    // Transactions
    for (String command : asList("DISCARD", "EXEC", "MULTI", "UNWATCH", "WATCH")) {
      sanitizers.put(command, KeepAllArgs.INSTANCE);
    }

    SANITIZERS = unmodifiableMap(sanitizers);
  }

  public static RedisCommandSanitizer create(boolean statementSanitizationEnabled) {
    return new RedisCommandSanitizer(statementSanitizationEnabled);
  }

  private final boolean statementSanitizationEnabled;

  private RedisCommandSanitizer(boolean statementSanitizationEnabled) {
    this.statementSanitizationEnabled = statementSanitizationEnabled;
  }

  public String sanitize(String command, List<?> args) {
    if (!statementSanitizationEnabled) {
      return KeepAllArgs.INSTANCE.sanitize(command, args);
    }
    return SANITIZERS
        .getOrDefault(command.toUpperCase(Locale.ROOT), DEFAULT)
        .sanitize(command, args);
  }

  interface CommandSanitizer {
    String sanitize(String command, List<?> args);
  }

  enum KeepAllArgs implements CommandSanitizer {
    INSTANCE;

    @Override
    public String sanitize(String command, List<?> args) {
      StringBuilder sanitized = new StringBuilder(command);
      for (Object arg : args) {
        sanitized.append(" ").append(argToString(arg));
      }
      return sanitized.toString();
    }
  }

  // keeps only a chosen number of arguments
  // example for num=2: CMD arg1 arg2 ? ?
  static final class CommandAndNumArgs implements CommandSanitizer {
    private final int numOfArgsToKeep;

    CommandAndNumArgs(int numOfArgsToKeep) {
      this.numOfArgsToKeep = numOfArgsToKeep;
    }

    @Override
    public String sanitize(String command, List<?> args) {
      StringBuilder sanitized = new StringBuilder(command);
      for (int i = 0; i < numOfArgsToKeep && i < args.size(); ++i) {
        sanitized.append(" ").append(argToString(args.get(i)));
      }
      for (int i = numOfArgsToKeep; i < args.size(); ++i) {
        sanitized.append(" ?");
      }
      return sanitized.toString();
    }
  }

  // keeps only chosen number of arguments and then every second one
  // example for num=2: CMD arg1 arg2 key1 ? key2 ?
  static final class MultiKeyValue implements CommandSanitizer {
    private final int numOfArgsBeforeKeyValue;

    MultiKeyValue(int numOfArgsBeforeKeyValue) {
      this.numOfArgsBeforeKeyValue = numOfArgsBeforeKeyValue;
    }

    @Override
    public String sanitize(String command, List<?> args) {
      StringBuilder sanitized = new StringBuilder(command);
      // append all "initial" arguments before key-value pairs start
      for (int i = 0; i < numOfArgsBeforeKeyValue && i < args.size(); ++i) {
        sanitized.append(" ").append(argToString(args.get(i)));
      }

      // loop over keys only
      for (int i = numOfArgsBeforeKeyValue; i < args.size(); i += 2) {
        sanitized.append(" ").append(argToString(args.get(i))).append(" ?");
      }
      return sanitized.toString();
    }
  }

  enum Eval implements CommandSanitizer {
    INSTANCE;

    @Override
    public String sanitize(String command, List<?> args) {
      StringBuilder sanitized = new StringBuilder(command);

      // get the number of keys passed from the command itself (second arg)
      int numberOfKeys = 0;
      if (args.size() > 2) {
        try {
          numberOfKeys = Integer.parseInt(argToString(args.get(1)));
        } catch (NumberFormatException ignored) {
          // Ignore
        }
      }

      int i = 0;
      // log the script, number of keys and all keys
      for (; i < (numberOfKeys + 2) && i < args.size(); ++i) {
        sanitized.append(" ").append(argToString(args.get(i)));
      }
      // mask the rest
      for (; i < args.size(); ++i) {
        sanitized.append(" ?");
      }
      return sanitized.toString();
    }
  }

  static String argToString(Object arg) {
    if (arg instanceof byte[]) {
      return new String((byte[]) arg, StandardCharsets.UTF_8);
    } else {
      return String.valueOf(arg);
    }
  }
}
