-- Sanitize the Lua environment to:
--   Prevent scripts from accessing files outside the game directory
--   Prevent usage of libraries other than


-- Utility: Convert a list of strings into a map of string->true
local function tmap(t)
  local u = {}; for _, n in ipairs(t) do u[n] = true end; return u;
end

-- Save banned global variables for wrapping with safe functions
local old_io = io
local old_require = require
local old_os = os
local old_debug = debug
local old_package = package
local old_allowffi = _bllua_allowffi

-- Remove all global variables except a whitelist
local ok_names = tmap {
  '_G', '_bllua_ts', '_bllua_on_unload', '_bllua_on_error',
  'string', 'table', 'math', 'coroutine', 'bit',
  'pairs', 'ipairs', 'next', 'unpack', 'select',
  'error', 'assert', 'pcall', 'xpcall',
  'type', 'tostring', 'tonumber',
  'loadstring',
  'getmetatable', 'setmetatable',
  'rawget', 'rawset', 'rawequal', 'rawlen',
  'module', '_VERSION',
}
local not_ok_names = {}
for n, _ in pairs(_G) do
  if not ok_names[n] then
    table.insert(not_ok_names, n)
  end
end
for _, n in ipairs(not_ok_names) do
  _G[n] = nil
end

-- Sanitize file paths to point only to allowed files within the game directory
-- List of allowed directories for reading/writing
-- modules/lualib is also allowed as read-only
local allowed_dirs = tmap {
  'add-ons', 'base', 'config', 'saves', 'screenshots', 'shaders'
}
-- List of disallowed file extensions - basically executable file extensions
-- Note that even without this protection, exploiting would still require somehow
--   getting a file within the allowed directories to autorun,
--   so this is just a precaution.
local disallowed_exts = tmap {
  -- windows
  'bat', 'bin', 'cab', 'cmd', 'com', 'cpl', 'ex_', 'exe', 'gadget', 'inf', 'ins', 'inx', 'isu',
  'job', 'jse', 'lnk', 'msc', 'msi', 'msp', 'mst', 'paf', 'pif', 'ps1', 'reg', 'rgs', 'scr',
  'sct', 'shb', 'shs', 'u3p', 'vb', 'vbe', 'vbs', 'vbscript', 'ws', 'wsf', 'wsh',
  -- linux
  'csh', 'ksh', 'out', 'run', 'sh',
  -- mac/other
  'action', 'apk', 'app', 'command', 'ipa', 'osx', 'prg', 'workflow',
}
-- Arguments: file name (relative to game directory), boolean true if only reading
-- Return: clean file path if allowed (or nil if disallowed),
--         error string (or nil if allowed)
local function safe_path(fn, readonly)
  if type(fn) ~= 'string' then
    return nil, 'Filename must be a string'
  end
  fn = fn:gsub('\\', '/')
  fn = fn:gsub('^ +', '')
  fn = fn:gsub(' +$', '')
  -- whitelist characters
  local ic = fn:find('[^a-zA-Z0-9_%-/ %.]')
  if ic then
    return nil, 'Filename \'' .. fn .. '\' contains invalid character \'' ..
        fn:sub(ic, ic) .. '\' at position ' .. ic
  end
  -- disallow up-dirs, absolute paths, and relative paths
  -- './' and '../' are possible in scripts, because they're processed into
  -- absolute paths in util.lua before reaching here
  if fn:find('^%.') or fn:find('%.%.') or fn:find(':') or fn:find('^/') then
    return nil, 'Filename \'' .. fn .. '\' contains invalid sequence'
  end
  -- allow only whitelisted dirs
  local dir = fn:match('^([^/]+)/')
  if not (dir and (
        allowed_dirs[dir:lower()] or
        (readonly and fn:find('^modules/lualib/'))))
  then
    return nil, 'File is in disallowed directory ' .. (dir or 'nil')
  end
  -- disallow blacklisted extensions
  local ext = fn:match('%.([^/%.]+)$')
  if ext and disallowed_exts[ext:lower()] then
    return nil, 'Filename \'' .. fn .. '\' has disallowed extension \'' ..
        (ext or '') .. '\''
  end
  return fn, nil
end

-- Wrap io.open with path sanitization
function _bllua_io_open(fn, md)
  md = md or 'r'
  local readonly = md == 'r' or md == 'rb'
  local fns, err = safe_path(fn, readonly)
  if fns then
    return old_io.open(fns, md)
  else
    return nil, err
  end
end

-- Allow io.type (works on file handles returned by io.open)
function _bllua_io_type(f)
  return old_io.type(f)
end

-- Wrap require with a blacklist for unsafe built-in modules
-- List of allowed packages to require
-- Note that util.lua wraps this and provides 'require',
--   only falling back here if the package is not found in user files
local disallowed_packages = tmap {
  'ffi', 'debug', 'package', 'io', 'os',
  '_bllua_ts',
}
if old_allowffi then disallowed_packages['ffi'] = nil end
function _bllua_requiresecure(name)
  if name:find('[^a-zA-Z0-9_%-%.]') or name:find('%.%.') or
      name:find('^%.') or name:find('%.$') then
    error('require: package name contains invalid character', 3)
  elseif disallowed_packages[name] then
    error('require: attempt to require disallowed module \'' .. name .. '\'', 3)
  else
    -- todo: reimplement require to not use package.* stuff?
    return old_require(name)
  end
end

package = {
  seeall = old_package.seeall,
}

-- Provide limited debug
debug = {
  traceback = old_debug.traceback,
  getinfo = old_debug.getinfo,
  getfilename = old_debug.getfilename, -- defined in lua.env.lua
}

print = _bllua_ts.echo
print('  Executed bllua-env-safe.lua')