#!/usr/bin/env ruby

require "bundler/inline"

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

  gem "declarative_policy", "~> 1.0"
  gem "erb", "~> 4.0"
  gem "globalid", "~> 1.0"
  gem "google-protobuf", "~> 3.0"
  gem "rack", "~> 3.0"
  gem "rackup", "~> 2.0"
  gem "saml-kit", "~> 1.0"
  gem "twirp", "~> 1.0"
  gem "webrick", "~> 1.0"
end

lib_path = Pathname.new(__FILE__).parent.parent.join('lib').realpath.to_s
$LOAD_PATH.unshift(lib_path) unless $LOAD_PATH.include?(lib_path)

require 'authx'

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

class JWT
  attr_reader :claims

  def initialize(claims)
    @claims = claims
  end

  def to_jwt
    [
      Base64.strict_encode64(JSON.generate({alg: "RS256", typ: "JWT"})),
      Base64.strict_encode64(JSON.generate(claims)),
      Base64.strict_encode64(JSON.generate({})),
    ].join(".")
  end
end

class User
  def initialize(attributes)
    @attributes = attributes
  end

  def name_id_for(name_id_format)
    if name_id_format == Saml::Kit::Namespaces::EMAIL_ADDRESS
      @attributes[:email]
    else
      @attributes[:id]
    end
  end

  def assertion_attributes_for(request)
    {
      custom: 'custom attribute',
      email: @attributes[:email],
      access_token: JWT.new(sub: SecureRandom.uuid, iat: Time.now.to_i).to_jwt,
    }
  end
end

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

class IdentityProvider
  def call(env)
    path = env['PATH_INFO']
    case env['REQUEST_METHOD']
    when 'GET'
      case path
      when '/.well-known/openid-configuration'
        return openid_metadata
      when '/.well-known/oauth-authorization-server'
        return oauth_metadata
      when '/.well-known/webfinger' # RFC-7033
        return not_found
      when "/saml/metadata.xml"
        return saml_metadata
      when "/saml/new"
        # TODO:: render a login page
        return saml_post_back(Rack::Request.new(env))
      when "/oauth/authorize" # RFC-6749
        return get_authorize(Rack::Request.new(env))
      else
        return not_found
      end
    when 'POST'
      case path
      when "/saml/new"
        return saml_post_back(Rack::Request.new(env))
      when "/oauth/authorize" # RFC-6749
        return post_authorize(Rack::Request.new(env))
      when "/oauth/token" # RFC-6749
        return [200, { 'Content-Type' => "application/json" }, [JSON.pretty_generate({
          access_token: JWT.new(sub: SecureRandom.uuid, iat: Time.now.to_i).to_jwt,
          token_type: "Bearer",
          expires_in: 3600,
          refresh_token: SecureRandom.hex(32)
        })]]
      when "/oauth/revoke" # RFC-7009
        return not_found
      else
        return not_found
      end
    end
    not_found
  end

  private

  def saml_metadata
    xml = Saml::Kit::Metadata.build_xml do |builder|
      builder.contact_email = 'hi@example.com'
      builder.organization_name = "Acme, Inc"
      builder.organization_url = "https://example.com"
      builder.build_identity_provider do |x|
        x.add_single_sign_on_service("#{$scheme}://#{$host}/saml/new", binding: :http_post)
        x.name_id_formats = [Saml::Kit::Namespaces::EMAIL_ADDRESS]
        x.attributes << :Username
      end
    end

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

  # GET /.well-known/oauth-authorization-server
  def oauth_metadata
    [200, { 'Content-Type' => "application/json" }, [JSON.pretty_generate({
      issuer: "#{$scheme}://#{$host}/.well-known/oauth-authorization-server",
      authorization_endpoint: "#{$scheme}://#{$host}/oauth/authorize",
      token_endpoint: "#{$scheme}://#{$host}/oauth/token",
      jwks_uri: "", # RFC-7517
      registration_endpoint: "", # RFC-7591
      scopes_supported: ["openid", "profile", "email"],
      response_types_supported: ["code", "code id_token", "id_token", "token id_token"],
      response_modes_supported: ["query", "fragment", "form_post"],
      grant_types_supported: ["authorization_code", "implicit"], # RFC-7591
      token_endpoint_auth_methods_supported: ["client_secret_basic"], # RFC-7591
      token_endpoint_auth_signing_alg_values_supported: ["RS256"],
      service_documentation: "",
      ui_locales_supported: ["en-US"],
      op_policy_uri: "",
      op_tos_uri: "",
      revocation_endpoint: "#{$scheme}://#{$host}/oauth/revoke", # RFC-7009
      revocation_endpoint_auth_methods_supported: ["client_secret_basic"],
      revocation_endpoint_auth_signing_alg_values_supported: ["RS256"],
      introspection_endpoint: "#{$scheme}://#{$host}/oauth/introspect", # RFC-7662
      introspection_endpoint_auth_methods_supported: ["client_secret_basic"],
      introspection_endpoint_auth_signing_alg_values_supported: ["RS256"],
      code_challenge_methods_supported: [], # RFC-7636
    })]]
  end

  def get_authorize(request)
    template = <<~ERB
      <!doctype html>
      <html>
        <head><title></title></head>
        <body>
          <h2>Authorize?</h2>
          <form action="/oauth/authorize" method="post">
            <input type="hidden" name="client_id" value="<%= request.params['client_id'] %>" />
            <input type="hidden" name="scope" value="<%= request.params['scope'] %>" />
            <input type="hidden" name="redirect_uri" value="<%= request.params['redirect_uri'] %>" />
            <input type="hidden" name="response_mode" value="<%= request.params['response_mode'] %>" />
            <input type="hidden" name="response_type" value="<%= request.params['response_type'] %>" />
            <input type="hidden" name="state" value="<%= request.params['state'] %>" />
            <input type="hidden" name="code_challenge_method" value="<%= request.params['code_challenge_method'] %>" />
            <input type="hidden" name="code_challenge" value="<%= request.params['code_challenge'] %>" />
            <input type="submit" value="Submit" />
          </form>
        </body>
      </html>
    ERB
    html = ERB.new(template, trim_mode: '-').result(binding)
    [200, { 'Content-Type' => "text/html" }, [html]]
  end

  def post_authorize(request)
    params = request.params.slice('client_id', 'redirect_uri', 'response_type', 'response_mode', 'state', 'code_challenge_method', 'code_challenge', 'scope')
    case params['response_type']
    when 'code'
      case params['response_mode']
      when 'fragment'
        return [302, { 'Location' => "#{params['redirect_uri']}#code=#{SecureRandom.uuid}&state=#{params['state']}" }, []]
      when 'query'
        return [302, { 'Location' => "#{params['redirect_uri']}?code=#{SecureRandom.uuid}&state=#{params['state']}" }, []]
      else
        # TODO:: form post
      end

    when 'token'
      return not_found
    else
      return not_found
    end
  end

  # GET /.well-known/openid-configuration
  def openid_metadata
    [200, { 'Content-Type' => "application/json" }, [JSON.pretty_generate({
      issuer: "#{$scheme}://#{$host}/.well-known/oauth-authorization-server",
      authorization_endpoint: "#{$scheme}://#{$host}/oauth/authorize",
      token_endpoint: "#{$scheme}://#{$host}/oauth/token",
      userinfo_endpoint: "#{$scheme}://#{$host}/oidc/user/",
      jwks_uri: "", # RFC-7517
      registration_endpoint: nil,
      scopes_supported: ["openid", "profile", "email"],
      response_types_supported: ["code", "code id_token", "id_token", "token id_token"],
      response_modes_supported: ["query", "fragment", "form_post"],
      grant_types_supported: ["authorization_code", "implicit"], # RFC-7591
      acr_values_supported: [],
      subject_types_supported: ["pairwise", "public"],
      id_token_signing_alg_values_supported: ["RS256"],
      id_token_encryption_alg_values_supported: [],
      id_token_encryption_enc_values_supported: [],
      userinfo_signing_alg_values_supported: ["RS256"],
      userinfo_encryption_alg_values_supported: [],
      userinfo_encryption_enc_values_supported: [],
      request_object_signing_alg_values_supported: ["none", "RS256"],
      request_object_encryption_alg_values_supported: [],
      request_object_encryption_enc_values_supported: [],
      token_endpoint_auth_methods_supported: ["client_secret_post", "client_secret_basic", "client_secret_jwt", "private_key_jwt"],
      token_endpoint_auth_signing_alg_values_supported: [],
      display_values_supported: [],
      claim_types_supported: ["normal", "aggregated", "distributed"],
      claims_supported: [
        "acr",
        "auth_time",
        "email",
        "email_verified",
        "family_name",
        "given_name",
        "iss",
        "locale",
        "name",
        "nickname",
        "picture",
        "profile",
        "sub",
        "website"
      ],
      service_documentation: nil,
      claims_locales_supported: [],
      ui_locales_supported: ["en-US"],
      claims_parameter_supported: false,
      request_parameter_supported: false,
      request_uri_paramater_supported: false,
      require_request_uri_registration: false,
      op_policy_uri: "",
      op_tos_uri: "",
    })]]
  end

  def saml_post_back(request)
    params = saml_params_from(request)
    saml_request = binding_for(request).deserialize(params)
    @builder = nil
    url, saml_params = saml_request.response_for(
      User.new({ id: SecureRandom.uuid, email: "example@example.com" }),
      binding: :http_post,
      relay_state: params[:RelayState]
    ) do |builder|
      builder.embed_signature = true
      @builder = builder
    end
    template = <<~ERB
      <!doctype html>
      <html>
        <head><title></title></head>
        <body>
          <h2>Recieved SAML Request</h2>
          <textarea readonly="readonly" disabled="disabled" cols=225 rows=6><%=- saml_request.to_xml(pretty: true) -%></textarea>

          <h2>Sending SAML Response (IdP -> SP)</h2>
          <textarea readonly="readonly" disabled="disabled" cols=225 rows=40><%=- @builder.build.to_xml(pretty: true) -%></textarea>
          <form action="<%= url %>" method="post">
            <%- saml_params.each do |(key, value)| -%>
              <input type="hidden" name="<%= key %>" value="<%= value %>" />
            <%- end -%>
            <input type="submit" value="Submit" />
          </form>
        </body>
      </html>
    ERB
    erb = ERB.new(template, trim_mode: '-')
    html = erb.result(binding)
    [200, { 'Content-Type' => "text/html" }, [html]]
  end

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

  def saml_params_from(request)
    if request.post?
      {
        "SAMLRequest" => request.params["SAMLRequest"],
        "RelayState" => request.params["RelayState"],
      }
    else
      query_string = request.query_string
      on = query_string.include?("&amp;") ? "&amp;" : "&"
      Hash[query_string.split(on).map { |x| x.split("=", 2) }].symbolize_keys
    end
  end

  def binding_for(request)
    location = "#{$scheme}://#{$host}/saml/new"
    if request.post?
      Saml::Kit::Bindings::HttpPost
        .new(location: location)
    else
      Saml::Kit::Bindings::HttpRedirect
        .new(location: location)
    end
  end
end

if __FILE__ == $0
  app = Rack::Builder.new do
    use Rack::CommonLogger
    use Rack::Reloader
    map "/twirp" do
      # https://github.com/arthurnn/twirp-ruby/wiki/Service-Handlers
      run ::Authx::Rpc::AbilityService.new(::Authx::Rpc::AbilityHandler.new)
    end

    run IdentityProvider.new
  end.to_app

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