hue.lua 1.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. local argparse = require("lib/argparse")
  2. local cookie = require("lib/cookie")
  3. local uri = require("lib/uri-f570bf7")
  4. local inspect = require("lib/inspect")
  5. local hue = {}
  6. hue.parse_args = function(args)
  7. local parser = argparse()
  8. parser:option "-X" "--method"
  9. :description "HTTP method"
  10. parser:option "-s" "--session"
  11. :description "session cookie"
  12. parser:flag "-v" "--verbose"
  13. :description "print verbosely"
  14. local args = parser:parse(args)
  15. if args["method"] then
  16. wrk.method = args["method"]
  17. end
  18. if args["session"] then
  19. wrk.headers["Cookie"] = cookie.build({
  20. sessionid=args["session"]
  21. })
  22. end
  23. hue.verbose = args.verbose
  24. end
  25. hue.handle_redirect = function(status, location)
  26. local location = uri.parse(location)
  27. wrk.scheme = location.scheme
  28. wrk.host = location.host
  29. wrk.port = location.port
  30. wrk.path = location.path
  31. if location.query then
  32. wrk.path = wrk.path .. "?" .. location.query
  33. end
  34. if location.fragment then
  35. wrk.path = wrk.path .. "#" .. location.fragment
  36. end
  37. if status == 303 then
  38. wrk.method = "GET"
  39. end
  40. end
  41. hue.set_cookies = function(set_cookie)
  42. local req_cookie = cookie.parse(wrk.headers["Cookie"] or "")
  43. local rep_cookie = cookie.parse(set_cookie)
  44. for key, val in pairs(rep_cookie) do
  45. -- ignore the cookie metadata
  46. if key ~= "Path" and key ~= "expires" and key ~= "HttpOnly" and key ~= "Max-Age" then
  47. req_cookie[key] = val
  48. end
  49. -- Make sure we handle the
  50. if key == "csrftoken" then
  51. wrk.headers["X-CSRFToken"] = val
  52. end
  53. end
  54. wrk.headers["Cookie"] = cookie.build(req_cookie)
  55. end
  56. return hue