TiagoVal97--cheatspecter/mock/assembler.mjs

// The pretend console's assembler.
//
// The real console assembles with Keystone (ps5debug-NG, assemble.c). When
// Keystone's Python module is installed on this computer the mock hands the
// text to the very same assembler, so what tests and the demo console
// produce is what the PS5 would produce:
//
//   pip install keystone-engine
//
// Without it, a small assembler written here stands in (see jsAssemble). It
// knows the everyday instructions, which is enough to try code cheats and the
// rule builder on the demo console, and says so when asked for anything more.

import { spawnSync } from 'node:child_process';
import { createRequire } from 'node:module';

const KS_SCRIPT = `
import sys, json
import keystone as K
req = json.load(sys.stdin)
ks = K.Ks(K.KS_ARCH_X86, K.KS_MODE_64)
try:
    enc, count = ks.asm(req["text"], int(req["base"]))
    print(json.dumps({"hex": bytes(enc or []).hex(), "count": count}))
except K.KsError as e:
    print(json.dumps({"error": str(e), "errno": e.errno}))
`;

let keystone = null;

/** True when Keystone can be used from here. Asked once; the answer is kept. */
export function hasKeystone() {
  if (keystone === null) {
    if (process.env.CS_NO_KEYSTONE) keystone = false;
    else {
      const r = spawnSync('python3', ['-c', 'import keystone'], { timeout: 20000 });
      keystone = r.status === 0;
    }
  }
  return keystone;
}

function keystoneAssemble(text, base) {
  const r = spawnSync('python3', ['-c', KS_SCRIPT], { input: JSON.stringify({ text, base: base.toString() }), timeout: 20000, encoding: 'utf8' });
  if (r.status !== 0) return { error: 'the assembler could not be started' };
  let out;
  try {
    out = JSON.parse(r.stdout);
  } catch {
    return { error: 'the assembler gave no answer' };
  }
  if (out.error) return { error: out.error, errno: out.errno };
  return { bytes: Buffer.from(out.hex, 'hex'), count: out.count };
}

const number = (tok) => {
  const t = tok.trim();
  if (/^-?0x[0-9a-f]+$/i.test(t) || /^-?\d+$/.test(t)) return BigInt(t);
  if (/^[0-9][0-9a-f]*h$/i.test(t)) return BigInt('0x' + t.slice(0, -1));
  return null;
};

/** The stand-in: two passes, every instruction a fixed size. */
function miniAssemble(text, base) {
  const items = [];
  const labels = new Map();
  let at = 0;
  let count = 0;
  for (const raw of text.split(/[\n;]/)) {
    let line = raw.replace(/(\/\/|#).*$/, '').trim();
    for (;;) {
      const m = /^([A-Za-z_.$][A-Za-z0-9_.$]*):/.exec(line);
      if (!m) break;
      labels.set(m[1], at);
      line = line.slice(m[0].length).trim();
      count++;
    }
    if (!line) continue;
    const [, op, rest = ''] = /^(\S+)\s*(.*)$/.exec(line);
    const word = op.toLowerCase();
    const args = rest.split(',').map((s) => s.trim()).filter(Boolean);
    let size;
    if (word === 'nop' || word === 'ret' || word === 'int3') size = 1;
    else if (word === 'jmp' && args.length === 1) size = 5;
    else if (word === 'db' || word === '.byte') size = args.length;
    else if (word === '.word') size = args.length * 2;
    else if (word === '.long' || word === '.int') size = args.length * 4;
    else if (word === '.quad') size = args.length * 8;
    else if (word === '.fill' || word === '.zero' || word === '.space') {
      const n = number(args[0] ?? '');
      if (n === null || n < 0n || n > 0x100000n) return { error: `invalid operand (${rest})` };
      size = Number(n) * (word === '.fill' ? Number(number(args[1] ?? '1') ?? 1n) : 1);
    } else if (word === '.balign' || word === '.align') {
      const n = Number(number(args[0] ?? '') ?? 0n);
      if (n <= 0 || n & (n - 1)) return { error: `invalid operand (${rest})` };
      size = (n - (at % n)) % n;
    } else return { error: `invalid mnemonic (${word})` };
    items.push({ word, args, at, size });
    at += size;
    count++;
  }
  const out = Buffer.alloc(at);
  const value = (tok) => {
    const n = number(tok);
    if (n !== null) return n;
    if (labels.has(tok)) return base + BigInt(labels.get(tok));
    return null;
  };
  for (const it of items) {
    const { word, args } = it;
    if (word === 'nop') out[it.at] = 0x90;
    else if (word === 'ret') out[it.at] = 0xc3;
    else if (word === 'int3') out[it.at] = 0xcc;
    else if (word === 'jmp') {
      const to = value(args[0]);
      if (to === null) return { error: `invalid operand (${args[0]})` };
      const rel = to - (base + BigInt(it.at) + 5n);
      if (rel < -0x80000000n || rel > 0x7fffffffn) return { error: 'invalid fixup' };
      out[it.at] = 0xe9;
      out.writeInt32LE(Number(rel), it.at + 1);
    } else if (word === '.fill') out.fill(Number((number(args[2] ?? '0') ?? 0n) & 0xffn), it.at, it.at + it.size);
    else if (word === '.balign' || word === '.align') out.fill(0x90, it.at, it.at + it.size);
    else if (word === '.zero' || word === '.space') out.fill(0, it.at, it.at + it.size);
    else {
      const width = word === '.quad' ? 8 : word === '.long' || word === '.int' ? 4 : word === '.word' ? 2 : 1;
      for (let i = 0; i < args.length; i++) {
        const v = value(args[i]);
        if (v === null) return { error: `invalid operand (${args[i]})` };
        if (width === 1 && (v < -128n || v > 255n)) return { error: `invalid operand (${args[i]})` };
        const o = it.at + i * width;
        if (width === 8) out.writeBigUInt64LE(BigInt.asUintN(64, v), o);
        else out.writeUIntLE(Number(BigInt.asUintN(width * 8, v)), o, width);
      }
    }
  }
  return { bytes: out, count };
}

// ---- a small assembler in JavaScript, for when Keystone is not installed -----------------------
//
// So that code cheats can be tried on the demo console on any computer. It
// encodes with iced-x86 (the library the app already uses to decode) and
// knows the everyday instructions code cheats are made of; anything else is
// refused by name. Jumps are always the long form, so its output is valid
// but not byte-for-byte what Keystone gives.

let icedLib = null;
const SIZE_WORDS = { byte: 1, word: 2, dword: 4, qword: 8, xmmword: 16 };
const JCC = { z: 'e', nz: 'ne', nae: 'b', c: 'b', nb: 'ae', nc: 'ae', na: 'be', nbe: 'a', nge: 'l', nl: 'ge', ng: 'le', nle: 'g' };
const ALU = ['add', 'or', 'adc', 'sbb', 'and', 'sub', 'xor', 'cmp', 'mov', 'test'];
const UNARY = ['inc', 'dec', 'neg', 'not'];
const SCALAR = { movss: 4, movsd: 8, addss: 4, subss: 4, mulss: 4, divss: 4, addsd: 8, subsd: 8, mulsd: 8, divsd: 8, comiss: 4, ucomiss: 4, comisd: 8, ucomisd: 8 };
const VECTOR = ['movdqu', 'movups', 'movaps', 'movdqa', 'xorps', 'pxor'];
const cap = (w) => w[0].toUpperCase() + w.slice(1);

class Refused extends Error {}

function jsAssemble(text, base) {
  icedLib ??= createRequire(import.meta.url)('iced-x86');
  const { Instruction, MemoryOperand, Register, Code, Encoder } = icedLib;
  const registers = new Map();
  for (const name of Object.keys(Register)) {
    if (/^(R[A-D]X|E[A-D]X|[A-D][XLH]|[RE]?[SD]I|[SD]IL|[RE]?[SB]P|[SB]PL|R(8|9|1[0-5])[DWL]?|XMM([0-9]|1[0-5]))$/.test(name)) {
      const size = /^XMM/.test(name) ? 16 : /^(R[A-DS0-9]|R[SD]I|RBP)/.test(name) && !/[DWL]$/.test(name) ? 8 : /^E|D$/.test(name) ? 4 : /^([A-D]X|[SD]I|[SB]P)$|W$/.test(name) ? 2 : 1;
      // iced spells the low byte of r8..r15 "R8L"; assemblers write "r8b".
      registers.set(name.toLowerCase().replace(/^(r\d+)l$/, '$1b'), { reg: Register[name], size, xmm: /^XMM/.test(name) });
    }
  }
  const code = (name) => {
    const c = Code[name];
    if (c === undefined) throw new Refused('invalid operand');
    return c;
  };

  // ---- reading ----
  const items = [];
  const labels = new Map();
  let count = 0;
  for (const raw of text.split(/[\n;]/)) {
    let line = raw.replace(/(\/\/|#).*$/, '').trim();
    for (;;) {
      const m = /^([A-Za-z_.$][A-Za-z0-9_.$]*):/.exec(line);
      if (!m) break;
      items.push({ label: m[1] });
      line = line.slice(m[0].length).trim();
      count++;
    }
    if (!line) continue;
    const [, op, rest = ''] = /^(\S+)\s*(.*)$/.exec(line);
    const args = [];
    let depth = 0, cur = '';
    for (const ch of rest) {
      if (ch === '[') depth++;
      if (ch === ']') depth--;
      if (ch === ',' && depth === 0) {
        args.push(cur.trim());
        cur = '';
      } else cur += ch;
    }
    if (cur.trim()) args.push(cur.trim());
    items.push({ op: op.toLowerCase(), args, rest });
    count++;
  }

  const value = (tok, at) => {
    const n = number(tok);
    if (n !== null) return n;
    if (labels.has(tok)) return base + BigInt(labels.get(tok));
    if (at === null && /^[A-Za-z_.$][A-Za-z0-9_.$]*$/.test(tok) && items.some((i) => i.label === tok)) return base; // first pass
    throw new Refused(`invalid operand (${tok})`);
  };

  // Set while an instruction is being built: where the next one starts, and whether an operand needs to know.
  let nextIp = 0n;
  let usesNext = false;

  /** One operand: { reg }, { imm }, or { mem: () => MemoryOperand, size } */
  const operand = (text, pass) => {
    const t = text.trim();
    const r = registers.get(t.toLowerCase());
    if (r) return { ...r, kind: 'reg' };
    const m = /^(?:(byte|word|dword|qword|xmmword)\s+ptr\s+)?\[(.*)\]$/i.exec(t);
    if (!m) return { kind: 'imm', imm: value(t, pass) };
    let baseReg = Register.None, index = Register.None, scale = 1, disp = 0n, rip = false, plain = true, named = false;
    for (const part of m[2].replace(/\s+/g, '').replace(/-/g, '+-').split('+').filter(Boolean)) {
      const neg = part.startsWith('-');
      const term = neg ? part.slice(1) : part;
      const scaled = /^([a-z0-9]+)\*([1248])$/i.exec(term);
      const reg = registers.get((scaled ? scaled[1] : term).toLowerCase());
      if (term.toLowerCase() === 'rip') rip = true;
      else if (reg && !neg) {
        plain = false;
        if (scaled || baseReg !== Register.None) {
          if (index !== Register.None) throw new Refused('invalid operand');
          index = reg.reg;
          scale = scaled ? Number(scaled[2]) : 1;
        } else baseReg = reg.reg;
      } else {
        if (number(term) === null) named = true;
        disp += (neg ? -1n : 1n) * value(term, pass);
      }
    }
    if (rip && !plain) throw new Refused('invalid operand');
    const size = m[1] ? SIZE_WORDS[m[1].toLowerCase()] : 0;
    // "[rip + label]" and a bare "[address]" both mean that address, reached relative to the instruction.
    // "[rip + 5]", with only a number, is that many bytes past the instruction, as Keystone reads it.
    const past = rip && !named;
    if (past) usesNext = true;
    const mem = () =>
      rip || plain
        ? new MemoryOperand(Register.RIP, Register.None, 1, BigInt.asIntN(64, past ? nextIp + disp : disp), 8, false, Register.None)
        : new MemoryOperand(baseReg, index, scale, disp, disp === 0n ? 0 : 1, false, Register.None);
    return { kind: 'mem', mem, size };
  };

  const fits = (n, bits) => n >= -(1n << BigInt(bits - 1)) && n < (1n << BigInt(bits - 1));
  const imm32 = (n, size) => Number(BigInt.asIntN(Math.min(size, 4) * 8, n));

  const build = (it, pass) => {
    const { op } = it;
    const o = it.args.map((a) => operand(a, pass));
    const [a, b] = o;
    const none = o.length === 0;
    if (none && op === 'nop') return Instruction.create(code('Nopd'));
    if (none && op === 'ret') return Instruction.create(code('Retnq'));
    if (none && op === 'int3') return Instruction.create(code('Int3'));
    if (none && (op === 'pushfq' || op === 'popfq')) return Instruction.create(code(cap(op)));
    if ((op === 'push' || op === 'pop') && o.length === 1 && a.kind === 'reg' && a.size === 8) return Instruction.createReg(code(`${cap(op)}_r64`), a.reg);
    const jcc = /^j([a-z]+)$/.exec(op);
    if ((op === 'jmp' || op === 'call' || jcc) && o.length === 1) {
      if (a.kind === 'imm') return Instruction.createBranch(code(op === 'jmp' || op === 'call' ? `${cap(op)}_rel32_64` : `J${JCC[jcc[1]] ?? jcc[1]}_rel32_64`), BigInt.asUintN(64, a.imm));
      if (jcc && op !== 'jmp') throw new Refused('invalid operand');
      if (a.kind === 'reg' && a.size === 8) return Instruction.createReg(code(`${cap(op)}_rm64`), a.reg);
      if (a.kind === 'mem') return Instruction.createMem(code(`${cap(op)}_rm64`), a.mem());
    }
    if (op === 'lea' && o.length === 2 && a.kind === 'reg' && b.kind === 'mem' && (a.size === 8 || a.size === 4)) return Instruction.createRegMem(code(`Lea_r${a.size * 8}_m`), a.reg, b.mem());
    if (op === 'movabs' && o.length === 2 && a.kind === 'reg' && a.size === 8 && b.kind === 'imm') return Instruction.createRegI64(code('Mov_r64_imm64'), a.reg, BigInt.asIntN(64, b.imm));
    if (UNARY.includes(op) && o.length === 1 && a.kind !== 'imm' && !a.xmm) {
      const s = a.size;
      if (!s) throw new Refused('invalid operand');
      return a.kind === 'reg' ? Instruction.createReg(code(`${cap(op)}_rm${s * 8}`), a.reg) : Instruction.createMem(code(`${cap(op)}_rm${s * 8}`), a.mem());
    }
    if (ALU.includes(op) && o.length === 2 && !a.xmm && !b.xmm && a.kind !== 'imm') {
      const M = cap(op);
      const s = a.kind === 'reg' ? a.size : b.kind === 'reg' ? b.size : a.size;
      if (!s || s > 8) throw new Refused('invalid operand');
      if ((a.kind === 'reg' && b.kind === 'reg' && a.size !== b.size) || (a.kind === 'mem' && b.kind === 'mem')) throw new Refused('invalid operand');
      if ((a.kind === 'mem' && a.size && b.kind === 'reg' && a.size !== b.size) || (b.kind === 'mem' && b.size && a.size !== b.size)) throw new Refused('invalid operand');
      const bits = s * 8;
      if (b.kind === 'reg') return a.kind === 'reg' ? Instruction.createRegReg(code(`${M}_rm${bits}_r${bits}`), a.reg, b.reg) : Instruction.createMemReg(code(`${M}_rm${bits}_r${bits}`), a.mem(), b.reg);
      if (b.kind === 'mem') {
        if (op === 'test') throw new Refused('invalid operand');
        return Instruction.createRegMem(code(`${M}_r${bits}_rm${bits}`), a.reg, b.mem());
      }
      const n = b.imm;
      if (op === 'mov' && a.kind === 'reg') {
        if (s === 8) return fits(n, 32) ? Instruction.createRegI32(code('Mov_rm64_imm32'), a.reg, Number(n)) : Instruction.createRegI64(code('Mov_r64_imm64'), a.reg, BigInt.asIntN(64, n));
        return Instruction.createRegI32(code(`Mov_r${bits}_imm${bits}`), a.reg, imm32(n, s));
      }
      if (s === 8 && !fits(n, 32)) throw new Refused('invalid operand');
      const short = op !== 'mov' && op !== 'test' && s > 1 && fits(n, 8);
      const name = `${M}_rm${bits}_imm${short ? 8 : Math.min(bits, 32)}`;
      return a.kind === 'reg' ? Instruction.createRegI32(code(name), a.reg, imm32(n, s)) : Instruction.createMemI32(code(name), a.mem(), imm32(n, s));
    }
    if (op in SCALAR && o.length === 2) {
      const bits = SCALAR[op] * 8;
      if (a.xmm && (b.xmm || b.kind === 'mem')) {
        const c = code(`${cap(op)}_xmm_xmmm${bits}`);
        return b.xmm ? Instruction.createRegReg(c, a.reg, b.reg) : Instruction.createRegMem(c, a.reg, b.mem());
      }
      if (a.kind === 'mem' && b.xmm && (op === 'movss' || op === 'movsd')) return Instruction.createMemReg(code(`${cap(op)}_xmmm${bits}_xmm`), a.mem(), b.reg);
    }
    if (VECTOR.includes(op) && o.length === 2) {
      if (a.xmm && (b.xmm || b.kind === 'mem')) {
        const c = code(`${cap(op)}_xmm_xmmm128`);
        return b.xmm ? Instruction.createRegReg(c, a.reg, b.reg) : Instruction.createRegMem(c, a.reg, b.mem());
      }
      if (a.kind === 'mem' && b.xmm && op.startsWith('mov')) return Instruction.createMemReg(code(`${cap(op)}_xmmm128_xmm`), a.mem(), b.reg);
    }
    throw new Refused(ALU.includes(op) || UNARY.includes(op) || op in SCALAR || VECTOR.includes(op) || ['lea', 'push', 'pop', 'jmp', 'call', 'movabs'].includes(op) || jcc ? 'invalid operand' : `invalid mnemonic (${op})`);
  };

  const DATA = { '.byte': 1, '.word': 2, '.short': 2, '.long': 4, '.int': 4, '.quad': 8 };
  const encoder = new Encoder(64);
  const emit = (pass) => {
    const chunks = [];
    let at = 0;
    for (const it of items) {
      if (it.label !== undefined) {
        if (pass === null) labels.set(it.label, at);
        continue;
      }
      let bytes;
      if (it.op in DATA) {
        const width = DATA[it.op];
        bytes = Buffer.alloc(it.args.length * width);
        it.args.forEach((arg, i) => {
          const v = value(arg, pass);
          if (width === 8) bytes.writeBigUInt64LE(BigInt.asUintN(64, v), i * 8);
          else bytes.writeUIntLE(Number(BigInt.asUintN(width * 8, v)), i * width, width);
        });
      } else if (it.op === '.fill' || it.op === '.zero' || it.op === '.space') {
        const n = Number(value(it.args[0] ?? '', pass)) * (it.op === '.fill' ? Number(value(it.args[1] ?? '1', pass)) : 1);
        if (!(n >= 0 && n <= 0x100000)) throw new Refused('invalid operand');
        bytes = Buffer.alloc(n, it.op === '.fill' ? Number(value(it.args[2] ?? '0', pass) & 0xffn) : 0);
      } else if (it.op === '.balign' || it.op === '.align') {
        const n = Number(value(it.args[0] ?? '', pass));
        if (n <= 0 || n & (n - 1)) throw new Refused('invalid operand');
        bytes = Buffer.alloc((n - (at % n)) % n, 0x90);
      } else if (it.op.startsWith('.')) throw new Refused(`unknown directive (${it.op})`);
      else {
        const ip = base + BigInt(at);
        const once = () => {
          const ins = build(it, pass);
          try {
            encoder.encode(ins, ip);
            return Buffer.from(encoder.takeBuffer());
          } catch {
            throw new Refused('invalid fixup');
          } finally {
            ins.free();
          }
        };
        usesNext = false;
        nextIp = ip;
        bytes = once();
        if (usesNext) {
          // Its length is the same wherever it points, so now the real end of the instruction is known.
          nextIp = ip + BigInt(bytes.length);
          bytes = once();
        }
      }
      chunks.push(bytes);
      at += bytes.length;
    }
    return Buffer.concat(chunks);
  };
  try {
    // Every instruction has one size whatever its addresses turn out to be, so
    // one pass places the labels and a second fills the addresses in.
    emit(null);
    return { bytes: emit(1), count };
  } catch (e) {
    if (e instanceof Refused) return { error: e.message };
    return { error: 'invalid operand' };
  } finally {
    encoder.free();
  }
}

export { jsAssemble };

/**
 * Assembles `text` as if placed at `base` (a bigint). `opts.keystone: false`
 * leaves Keystone out even when it is installed.
 * Returns { bytes, count } or { error }.
 */
export function assemble(text, base, opts = {}) {
  // The stand-in first: it is instant, and it keeps accepting "db", which older tests use.
  const mini = miniAssemble(text, base);
  if (!mini.error) return mini;
  if (opts.keystone !== false && hasKeystone()) return keystoneAssemble(text, base);
  const js = jsAssemble(text, base);
  if (js.error && /invalid mnemonic/.test(js.error)) {
    return { error: `${js.error}: this pretend console only knows the everyday instructions (mov, cmp, add, jumps, calls, push and pop, and the common decimal-number ones). A PS5 assembles any instruction` };
  }
  return js;
}