gem 'libsaml', require: 'saml'

require 'openssl'
require 'base64'

pem = File.open("config/ssl/trust-federate.cert").read
cert = OpenSSL::X509::Certificate.new(pem)
output = Base64.encode64(cert.to_der).gsub("\n", "")

Saml.setup do |config|
  config.register_store :file, Saml::ProviderStores::File.new("config/metadata", "config/ssl/key.pem"), default: true
end

class SamlProvider < ActiveRecord::Base
  include Saml::Provider

  def self.find_by_entity_id(entity_id)
    find_by entity_id: entity_id
  end
end

class SamlController < ApplicationController
  extend Saml::Rails::ControllerHelper
  current_provider "entity_id"

  def request_authentication
    provider = Saml.provider("my:very:original:entityid")
    destination = provider.single_sign_on_service_url(Saml::ProtocolBinding::HTTP_POST)

    authn_request = Saml::AuthnRequest.new(:destination => destination)

    session[:authn_request_id] = authn_request._id

    @saml_attributes = Saml::Bindings::HTTPPost.create_form_attributes(authn_request)

    render text: @saml_attributes.to_yaml
  end

  def receive_response
    if params["SAMLart"]
      # provider should be of type Saml::Provider
      @response = Saml::Bindings::HTTPArtifact.resolve(request, provider.artifact_resolution_service_url)
    elsif params["SAMLResponse"]
      @response = Saml::Bindings::HTTPPost.receive_message(request, :response)
    else
       # handle invalid request
    end

    if @response && @response.success?
      if session[:authn_request_id] == @response.in_response_to
        @response.assertion.fetch_attribute('any_attribute')
      else
        # handle unrecognized response
      end
      reset_session # It's good practice to reset sessions after authenticating to mitigate session fixation attacks
    else
      # handle failure
    end
  end
end

  get "/saml/request_authentication" => "saml#request_authentication"
  get "/saml/receive_response" => "saml#receive_response"
  post "/saml/receive_response" => "saml#receive_response"

class SamlController < ActionController::Base
  extend Saml::Rails::ControllerHelper
  current_provider "entity_id"

  def receive_authn_request
    authn_request = if request.get?
      Saml::Bindings::HTTPRedirect.receive_message(request, type: :authn_request)
    elsif request.post?
      Saml::Bindings::HTTPPost.receive_message(request, :authn_request)
    else
      return head :not_allowed
    end
    request_id = authn_request._id

    session[:saml_request] = {
      request_id:    request_id,
      relay_state:   params['RelayState'],
      authn_request: authn_request.to_xml
    }

    if authn_request.invalid?
      redirect_to send_response_path(request_id: request_id)
    else
      redirect_to sign_in_path(return_to: send_response_path(request_id: request_id))
    end
  end

  def send_response
    return head :not_found if session[:saml_request][:request_id] != params[:request_id]

    authn_request = Saml::AuthnRequest.parse(session[:saml_request][:authn_request], single: true)

    response = if authn_request.invalid?
      build_failure(Saml::TopLevelCodes::REQUESTER, Saml::SubStatusCodes::REQUEST_DENIED)
    elsif account_signed_in?
      build_success_response
    else
      build_failure(Saml::TopLevelCodes::RESPONDER, Saml::SubStatusCodes::NO_AUTHN_CONTEXT, 'cancelled')
    end

    if authn_request.protocol_binding == Saml::ProtocolBinding::HTTP_POST
      # render an auto submit form with hidden fields set in the attributes hash
      @attribute = Saml::Bindings::HTTPPost.create_form_attributes(response, relay_state: session[:saml_request][:relay_state])
    else
      # handle unsupported binding
    end
  end

  private

  def build_failure(status_value, sub_status_value, status_detail)
    Saml::Response.new(in_response_to:   session[:saml_request][:request_id],
                       status_value:     status_value,
                       sub_status_value: sub_status_value,
                       status_detail:    status_detail)
  end

  def build_success_response(authn_request)
    assertion = Saml::Assertion.new(
      name_id:                 current_account.username, # Return anything that you can link to an account
      name_id_format:          'urn:oasis:names:tc:SAML:2.0:nameid-format:persistent',
      authn_context_class_ref: Saml::ClassRefs::PASSWORD_PROTECTED,
      in_response_to:          authn_request._id,
      recipient:               authn_request.assertion_url,
      audience:                authn_request.issuer)

    # adding custom attributes
    assertion.add_attribute('name', 'value')

    Saml::Response.new(in_response_to: authn_request._id,
                       assertion:      assertion,
                       status_value:   Saml::TopLevelCodes::SUCCESS)
  end
end