#!/usr/bin/env ruby

require "bundler/inline"

gemfile do
  source "https://rubygems.org"

  gem "base64", "~> 0.1"
  gem "erb", "~> 4.0"
  gem "net-hippie", "~> 1.0"
  gem "rack", "~> 3.0"
  gem "rack-session", "~> 2.0"
  gem "rackup", "~> 2.0"
  gem "saml-kit", "~> 1.0"
  gem "webrick", "~> 1.0"
end

$scheme = ENV.fetch("SCHEME", "http")
$port = ENV.fetch("PORT", 8283).to_i
$host = ENV.fetch("HOST", "localhost:#{$port}")
$idp_host = ENV.fetch("IDP_HOST", "localhost:8282")

Net::Hippie.logger = Logger.new($stdout, level: :debug)

class OnDemandRegistry < Saml::Kit::DefaultRegistry
  def metadata_for(entity_id)
    found = super(entity_id)
    return found if found

    register_url(entity_id, verify_ssl: false)
    super(entity_id)
  end
end

Saml::Kit.configure do |x|
  x.entity_id = "#{$scheme}://#{$host}/saml/metadata.xml"
  x.registry = OnDemandRegistry.new
  x.logger = Logger.new("/dev/stderr")
end

module OAuth
  class Client
    attr_reader :client_id, :client_secret, :http, :authz_host

    def initialize(authz_host, client_id, client_secret)
      @authz_host = authz_host
      @client_id = client_id
      @client_secret = client_secret
      @http = Net::Hippie::Client.new(headers: ::Net::Hippie::Client::DEFAULT_HEADERS.merge({
        'Authorization' => Net::Hippie.basic_auth(client_id, client_secret),
      }))
    end

    def [](key)
      server_metadata.fetch(key)
    end

    def authorize_uri(state: SecureRandom.uuid, response_type: "code", response_mode: "query", scope: "openid")
      [
        self[:authorization_endpoint],
        to_query(
          client_id: client_id,
          state: state,
          redirect_uri: redirect_uri,
          response_mode: response_mode,
          response_type: response_type,
          scope: scope,
        )
      ].join("?")
    end

    def exchange(grant_type:, code:, code_verifier: "not_implemented")
      with_http do |client|
        client.post(self[:token_endpoint], body: {
          grant_type: grant_type,
          code: code,
          code_verifier: code_verifier,
        })
      end
    end

    private

    def to_query(params = {})
      params.map do |(key, value)|
        [key, value].join("=")
      end.join("&")
    end

    def redirect_uri
      "#{$scheme}://#{$host}/oauth/callback"
    end

    def with_http
      http.with_retry do |client|
        yield client
      end
    end

    def server_metadata
      @server_metadata ||=
        with_http do |client|
          response = client.get("http://#{authz_host}/.well-known/oauth-authorization-server")
          JSON.parse(response.body, symbolize_names: true)
        end
    end
  end
end

module HTTPHelpers
  def not_found
    [404, { 'X-Backend-Server' => 'UI' }, []]
  end

  def redirect_to(location)
    if location.start_with?("http")
      [302, { 'Location' => location }, []]
    else
      [302, { 'Location' => "http://ui.example.com:8080#{location}" }, []]
    end
  end

end

class UI
  include ::HTTPHelpers

  attr_reader :oauth_client

  def initialize(oauth_client)
    @oauth_client = oauth_client
  end

  def call(env)
    path = env['PATH_INFO']
    case env['REQUEST_METHOD']
    when 'GET'
      case path
      when "/oauth/callback"
        return oauth_callback(Rack::Request.new(env))
      when "/oidc/new"
        return redirect_to(oauth_client.authorize_uri)
      when "/saml/metadata.xml"
        return metadata
      when "/saml/new"
        return saml_post_to_idp(Rack::Request.new(env))
      else
        return redirect_to("/saml/new")
      end
    when 'POST'
      case path
      when "/saml/assertions"
        return saml_assertions(Rack::Request.new(env))
      else
        return not_found
      end
    end
    not_found
  end

  private

  def metadata
    xml = Saml::Kit::Metadata.build_xml do |builder|
      builder.embed_signature = false
      builder.contact_email = 'ui@example.com'
      builder.organization_name = "Acme, Inc"
      builder.organization_url = "https://example.com"
      builder.build_service_provider do |x|
        x.name_id_formats = [Saml::Kit::Namespaces::PERSISTENT]
        x.add_assertion_consumer_service("#{$scheme}://#{$host}/saml/assertions", binding: :http_post)
      end
    end

    [200, { 'Content-Type' => "application/samlmetadata+xml" }, [xml]]
  end

  def oauth_callback(request)
    response = oauth_client.exchange(grant_type: "authorization_code", code: request.params['code'])
    [response.code, response.header, [response.body]]
  end

  def saml_post_to_idp(request)
    idp = Saml::Kit.registry.metadata_for("http://#{$idp_host}/saml/metadata.xml")
    relay_state = Base64.strict_encode64(JSON.generate(redirect_to: '/dashboard'))

    @saml_builder = nil
    uri, saml_params = idp.login_request_for(binding: :http_post, relay_state: relay_state) do |builder|
      @saml_builder = builder
    end

    template = <<~ERB
      <!doctype html>
      <html>
        <head><title></title></head>
        <body style="background-color: pink;">
          <h2>Sending SAML Request (SP -> IdP)</h2>
          <textarea readonly="readonly" disabled="disabled" cols=225 rows=6><%=- @saml_builder.to_xml(pretty: true) -%></textarea>

          <form id="idp-form" action="<%= uri %>" method="post">
            <%- saml_params.each do |(key, value)| -%>
              <input type="hidden" name="<%= key %>" value="<%= value %>" />
            <%- end -%>
            <input id="submit-button" type="submit" value="Submit" />
          </form>
        </body>
      </html>
    ERB
    html = ERB.new(template, trim_mode: '-').result(binding)
    [200, { 'Content-Type' => "text/html" }, [html]]
  end

  def saml_assertions(request)
    sp = Saml::Kit.registry.metadata_for("#{$scheme}://#{$host}/saml/metadata.xml")
    saml_binding = sp.assertion_consumer_service_for(binding: :http_post)
    saml_response = saml_binding.deserialize(request.params)
    raise saml_response.errors unless saml_response.valid?

    request.session[:access_token] = saml_response.attributes[:access_token]
    template = <<~ERB
      <!doctype html>
      <html>
        <head>
          <title></title>
        </head>
        <body style="background-color: pink;">
          <h2>Received SAML Response</h2>
          <textarea readonly="readonly" disabled="disabled" cols=220 rows=40><%=- saml_response.to_xml(pretty: true) -%></textarea>
        </body>
      </html>
    ERB
    html = ERB.new(template, trim_mode: '-').result(binding)
    [200, { 'Content-Type' => "text/html" }, [html]]
  end
end

if __FILE__ == $0
  app = Rack::Builder.new do
    use Rack::CommonLogger
    use Rack::Reloader
    use Rack::Session::Cookie, { domain: $host.split(":", 2)[0], path: "/", secret: SecureRandom.hex(64) }

    run UI.new(::OAuth::Client.new($idp_host, 'client_id', 'client_secret'))
  end.to_app

  Rackup::Server.start(app: app, Port: $port)
end