/**
 * @license
 * Copyright 2016 Google Inc. All rights reserved.
 * 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.
 */

package com.google.security.wycheproof;

import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.security.GeneralSecurityException;
import java.util.Random;

/**
 * A collection of utilities for testing random number generators. So far this util simply checks
 * that random numbers are not generated by java.util.Random. Eventually we plan to add detection
 * for other random number generators too.
 *
 * @author bleichen@google.com (Daniel Bleichenbacher)
 */
public class RandomUtil {
  // Constants for java.util.Random;
  static final long A = 0x5DEECE66DL;
  static final long A_INVERSE = 246154705703781L;
  static final long C = 0xBL;

  /** Given a state of a java.util.Random object compute the next state. */
  protected static long nextState(long seed) {
    return (seed * A + C) & ((1L << 48) - 1);
  }

  /** Give the state after stepping java.util.Random n times. */
  protected static long step(long seed, long n) {
    long a = A;
    long c = C;
    n = n & 0xffffffffffffL;
    while (n != 0) {
      if ((n & 1) == 1) {
        seed = seed * a + c;
      }
      c = c * (a + 1);
      a = a * a;
      n = n >> 1;
    }
    return seed & 0xffffffffffffL;
  }

  /** Given a state of a java.util.Random object compute the previous state. */
  protected static long previousState(long seed) {
    return ((seed - C) * A_INVERSE) & ((1L << 48) - 1);
  }

  /** Computes a seed that would initialize a java.util.Random object with a given state. */
  protected static long getSeedForState(long seed) {
    return seed ^ A;
  }

  protected static long getStateForSeed(long seed) {
    return (seed ^ A) & 0xffffffffffffL;
  }

  /**
   * Given two subsequent outputs x0 and x1 from java.util.Random this function computes the
   * internal state of java.util.Random after returning x0 or returns -1 if no such state exists.
   */
  protected static long getState(int x0, int x1) {
    long mask = (1L << 48) - 1;
    long multiplier = A;
    // The state of the random number generator after returning x0 is
    // l0 + eps for some 0 <= eps < 2**16.
    long l0 = ((long) x0 << 16) & mask;
    // The state of the random number generator after returning x1 is
    // l1 + delta for some 0 <= delta < 2**16.
    long l1 = ((long) x1 << 16) & mask;
    // We have l1 + delta = (l0 + eps)*multiplier + 0xBL (mod 2**48).
    // This allows to find an upper bound w for eps * multiplier mod 2**48
    // by assuming delta = 2**16-1.
    long w = (l1 - l0 * multiplier + 65535L - 0xBL) & mask;
    // The reduction eps * multiplier mod 2**48 only cuts off at most 3 bits.
    // Hence a simple search is sufficient. The maximal number of loops is 6.
    for (long em = w; em < (multiplier << 16); em += 1L << 48) {
      // If the high order bits of em are guessed correctly then
      // em == eps * multiplier + 65535 - delta.
      long eps = em / multiplier;
      long state0 = l0 + eps;
      long state1 = nextState(state0);
      if ((state1 & 0xffffffff0000L) == l1) {
        return state0;
      }
    }
    return -1;
  }

  /**
   * Find a seed such that this integer is the result of
   *
   * <pre>{@code
   * Random rand = new Random();
   * rand.setSeed(seed);
   * return new BigInteger(k, rand);
   * }</pre>
   *
   * where k is max(64, x.BitLength());
   *
   * <p>Returns -1 if no such seed exists.
   */
  // TODO(bleichen): We want to detect cases where some of the bits
  //   (i.e. most significant bits or least significant bits have
  //   been modified. Often this happens during the generation
  //   of primes or other things.
  // TODO(bleichen): This method is incomplete.
  protected static long getSeedFor(java.math.BigInteger x) {
    byte[] bytes = x.toByteArray();
    if (bytes.length == 0) {
      return -1;
    }
    ByteBuffer buffer = ByteBuffer.allocate(8);
    int offset = bytes[0] == 0 ? 1 : 0;
    if (bytes.length - offset < 8) {
      int size = bytes.length - offset;
      buffer.position(8 - size);
      buffer.put(bytes, offset, size);
    } else {
      buffer.put(bytes, offset, 8);
    }
    buffer.flip();
    buffer.order(java.nio.ByteOrder.LITTLE_ENDIAN);
    int x0 = buffer.getInt();
    int x1 = buffer.getInt();
    long state = getState(x0, x1);
    if (state == -1) {
      return -1;
    }
    return getSeedForState(previousState(state));
  }

  /** Attempts to find a seed such that it generates the prime p. Returns -1 if no seed is found. */
  static long getSeedForPrime(BigInteger p) {
    int confidence = 64;
    Random rand = new Random();
    int size = p.bitLength();
    // Prime generation often sets the most significant bit.
    // Hence, clearing the most significant bit can help to find
    // the seed used for the prime generation.
    for (BigInteger x : new BigInteger[] {p, p.clearBit(size - 1)}) {
      long seed = getSeedFor(x);
      if (seed != -1) {
        rand.setSeed(seed);
        BigInteger q = new BigInteger(size, confidence, rand);
        if (q.equals(p)) {
          return seed;
        }
      }
    }
    return -1;
  }

  /**
   * Checks whether p is a random prime. A prime generated with a secure random number generator
   * passes with probability > 1-2^{-32}. No checks are performed for primes smaller than 96 bits.
   *
   * @throws GeneralSecurityException if the prime was generated using java.util.Random
   */
  static void checkPrime(BigInteger p) throws GeneralSecurityException {
    // We can't reliably detect java.util.Random for small primes.
    if (p.bitLength() < 96) {
      return;
    }
    long seed = getSeedForPrime(p);
    if (seed != -1) {
      throw new GeneralSecurityException(
          "java.util.Random with seed " + seed + " was likely used to generate prime");
    }
  }
}