crypto.lua

-- Cryptographic utilities for diffusion.nvim
-- Provides secure random generation and WebSocket key handling

local Crypto = {}

-- WebSocket GUID constant for handshake
Crypto.WEBSOCKET_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

-- Generate a cryptographically secure UUID v4
function Crypto.generate_uuid()
  local template = "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx"
  local uuid = string.gsub(template, "[xy]", function(c)
    local v = (c == "x") and math.random(0, 0xf) or math.random(8, 0xb)
    return string.format("%x", v)
  end)
  return uuid
end

-- Generate secure random bytes
function Crypto.random_bytes(count)
  local bytes = {}
  for i = 1, count do
    table.insert(bytes, math.random(0, 255))
  end
  return bytes
end

-- Base64 encoding
function Crypto.base64_encode(data)
  local chars = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
  local result = {}
  
  -- Pad data to multiple of 3
  local padding = 3 - (#data % 3)
  if padding ~= 3 then
    for i = 1, padding do
      data = data .. '\0'
    end
  end
  
  -- Encode in groups of 3 bytes
  for i = 1, #data, 3 do
    local b1, b2, b3 = string.byte(data, i, i + 2)
    -- avoid bitwise operators for LuaJIT 5.1 compatibility
    local n = b1 * 65536 + b2 * 256 + b3
    
    local c1 = math.floor(n / 262144) % 64 -- 2^18
    local c2 = math.floor(n / 4096) % 64   -- 2^12
    local c3 = math.floor(n / 64) % 64     -- 2^6
    local c4 = n % 64
    
    table.insert(result, chars:sub(c1 + 1, c1 + 1))
    table.insert(result, chars:sub(c2 + 1, c2 + 1))
    table.insert(result, chars:sub(c3 + 1, c3 + 1))
    table.insert(result, chars:sub(c4 + 1, c4 + 1))
  end
  
  -- Apply padding
  if padding == 1 then
    result[#result] = '='
  elseif padding == 2 then
    result[#result] = '='
    result[#result - 1] = '='
  end
  
  return table.concat(result)
end

-- Base64 decoding
function Crypto.base64_decode(data)
  local chars = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
  local lookup = {}
  
  -- Build lookup table
  for i = 1, #chars do
    lookup[chars:sub(i, i)] = i - 1
  end
  
  -- Remove padding
  data = data:gsub('=', '')
  
  local result = {}
  local padding = (4 - (#data % 4)) % 4
  
  -- Decode in groups of 4 characters
  for i = 1, #data, 4 do
    local c1 = lookup[data:sub(i, i)] or 0
    local c2 = lookup[data:sub(i + 1, i + 1)] or 0
    local c3 = lookup[data:sub(i + 2, i + 2)] or 0
    local c4 = lookup[data:sub(i + 3, i + 3)] or 0
    
    local n = c1 * 262144 + c2 * 4096 + c3 * 64 + c4
    
    local b1 = math.floor(n / 65536) % 256
    local b2 = math.floor(n / 256) % 256
    local b3 = n % 256
    
    table.insert(result, string.char(b1))
    if i + 1 <= #data then
      table.insert(result, string.char(b2))
    end
    if i + 2 <= #data then
      table.insert(result, string.char(b3))
    end
  end
  
  return table.concat(result)
end

-- SHA-1 implementation in pure Lua
function Crypto.sha1(data)
  local ok, bit = pcall(require, 'bit')
  if not ok then
    bit = _G.bit or _G.bit32
  end
  assert(bit and bit.band and bit.bor and bit.bxor and bit.bnot and bit.lshift and bit.rshift, 'bit library not available')

  local function left_rotate(n, b)
    -- rotate-left using bit ops
    return bit.band(bit.bor(bit.lshift(n, b), bit.rshift(n, 32 - b)), 0xFFFFFFFF)
  end
  
  local function to_bytes(str)
    local bytes = {}
    for i = 1, #str do
      table.insert(bytes, string.byte(str, i))
    end
    return bytes
  end
  
  local function from_bytes(bytes)
    local result = {}
    for i = 1, #bytes do
      table.insert(result, string.char(bytes[i]))
    end
    return table.concat(result)
  end
  
  -- Convert string to bytes
  local message = to_bytes(data)
  local original_length = #message * 8
  
  -- Append single '1' bit (as 0x80 byte)
  table.insert(message, 0x80)
  
  -- Append zeros until length ≡ 448 (mod 512)
  while (#message % 64) ~= 56 do
    table.insert(message, 0)
  end
  
  -- Append original length as 64-bit big-endian (Lua 5.1-compatible)
  for i = 7, 0, -1 do
    local byte = math.floor(original_length / (2^(i * 8))) % 256
    table.insert(message, byte)
  end
  
  -- Initialize hash values
  local h0 = 0x67452301
  local h1 = 0xEFCDAB89
  local h2 = 0x98BADCFE
  local h3 = 0x10325476
  local h4 = 0xC3D2E1F0
  
  -- Process message in 512-bit chunks
  for chunk_start = 1, #message, 64 do
    local w = {}
    
    -- Break chunk into sixteen 32-bit big-endian words
    for i = 0, 15 do
      local byte_offset = chunk_start + i * 4
      local a = (message[byte_offset] or 0) * 16777216
      local b = (message[byte_offset + 1] or 0) * 65536
      local c = (message[byte_offset + 2] or 0) * 256
      local d = (message[byte_offset + 3] or 0)
      w[i] = (a + b + c + d) % 4294967296
    end
    
    -- Extend words
    for i = 16, 79 do
      local x = (w[i - 3] or 0)
      x = bit.bxor(x, w[i - 8] or 0)
      x = bit.bxor(x, w[i - 14] or 0)
      x = bit.bxor(x, w[i - 16] or 0)
      w[i] = left_rotate(x % 4294967296, 1)
    end
    
    -- Initialize working variables
    local a, b, c, d, e = h0, h1, h2, h3, h4
    
    -- Main loop
    for i = 0, 79 do
      local f, k
      
      if i <= 19 then
        f = bit.bor(bit.band(b, c), bit.band(bit.bnot(b), d))
        k = 0x5A827999
      elseif i <= 39 then
        f = bit.bxor(bit.bxor(b, c), d)
        k = 0x6ED9EBA1
      elseif i <= 59 then
        f = bit.bor(bit.band(b, c), bit.bor(bit.band(b, d), bit.band(c, d)))
        k = 0x8F1BBCDC
      else
        f = bit.bxor(bit.bxor(b, c), d)
        k = 0xCA62C1D6
      end
      
      local temp = (left_rotate(a, 5) + f + e + k + (w[i] or 0)) % 4294967296
      e = d
      d = c
      c = left_rotate(b, 30)
      b = a
      a = temp
    end
    
    -- Add chunk hash to result
    h0 = (h0 + a) % 4294967296
    h1 = (h1 + b) % 4294967296
    h2 = (h2 + c) % 4294967296
    h3 = (h3 + d) % 4294967296
    h4 = (h4 + e) % 4294967296
  end
  
  -- Convert hash to binary string
  local result = {}
  for _, h in ipairs({h0, h1, h2, h3, h4}) do
    for i = 3, 0, -1 do
      local byte = math.floor(h / (2^(i*8))) % 256
      table.insert(result, byte)
    end
  end
  
  return from_bytes(result)
end

-- Generate WebSocket accept key from client key
function Crypto.generate_websocket_accept(client_key)
  local concat_key = client_key .. Crypto.WEBSOCKET_GUID
  local hash = Crypto.sha1(concat_key)
  return Crypto.base64_encode(hash)
end

-- Validate UUID format
function Crypto.validate_uuid(uuid)
  if type(uuid) ~= "string" then
    return false
  end
  
  local pattern = "^[0-9a-f]{8}%-[0-9a-f]{4}%-[0-9a-f]{4}%-[0-9a-f]{4}%-[0-9a-f]{12}$"
  return uuid:match(pattern) ~= nil
end

-- Generate secure token with optional timestamp
function Crypto.generate_secure_token(include_timestamp)
  local uuid = Crypto.generate_uuid()
  
  if include_timestamp then
    local timestamp = string.format("%x", os.time())
    -- Replace last part of UUID with timestamp
    uuid = uuid:gsub("-([0-9a-f]+)$", "-" .. timestamp)
  end
  
  return uuid
end

-- Validate token format and check expiration if applicable
function Crypto.validate_token(token, max_age)
  if not Crypto.validate_uuid(token) then
    return false, "Invalid token format"
  end
  
  if max_age then
    -- Extract timestamp from last part if present
    local timestamp_hex = token:match("-([0-9a-f]+)$")
    if timestamp_hex and #timestamp_hex == 8 then
      local timestamp = tonumber(timestamp_hex, 16)
      if timestamp and (os.time() - timestamp) > max_age then
        return false, "Token expired"
      end
    end
  end
  
  return true
end

-- Initialize random seed
math.randomseed(os.time())

return Crypto