=begin
    Copyright 2010-2014 Tasos Laskos <tasos.laskos@gmail.com>

    Licensed under the Apache License, Version 2.0 (the "License");
    you may not use this file except in compliance with the License.
    You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
=end

module Arachni

#
# Session management class.
#
# Handles logins, provided log-out detection, stores and executes login sequences
# and provided general webapp session related helpers.
#
# @author Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
#
class Session
    include UI::Output
    include Utilities

    #
    # {Session} error namespace.
    #
    # All {Session} errors inherit from and live under it.
    #
    # @author Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
    #
    class Error < Arachni::Error

        #
        # Raised when a login check is required to perform an action but none
        # has been configured.
        #
        # @author Tasos "Zapotek" Laskos <tasos.laskos@gmail.com>
        #
        class NoLoginCheck < Error
        end
    end

    LOGIN_TRIES      = 5
    LOGIN_RETRY_WAIT = 5

    # @return   [Options]  options
    attr_reader :opts

    #
    # A block used to login to the webapp.
    #
    # The block should log the framework into the webapp and return `true` on
    # success, `false` on failure.
    #
    # @return   [Block]
    #
    attr_accessor :login_sequence

    #
    # A block used to check whether or not we're logged in to the webapp.
    #
    # It should:
    #
    #   * return `true` on success, `false` on failure.
    #   * expect 2 parameters, the first one being a hash of HTTP options and
    #       the second one an optional block.
    #
    # If a block has been set, the check should work async and pass the result
    # to the block, otherwise it should simply return the result.
    #
    # The result of the check should be `true` or `false`.
    #
    # A good example of this can be found in {#set_login_check}.
    #
    # @return   [Block]
    #
    # @see #set_login_check
    #
    attr_accessor :login_check

    #
    # Sets a login form and generates a login sequence from it.
    #
    # The form must be kosher, best be generated by one of the {Arachni::Element::Form}
    # helpers, {Parser} or {#find_login_form}.
    #
    # Once you get the right form you need to update it with the appropriate values
    # before passing it to this accessor.
    #
    # @return   [Element::Form]
    #
    attr_accessor :login_form

    def initialize( opts = Arachni::Options.instance )
        @opts = opts
    end

    # @return   [Array<Element::Cookie>]    session cookies
    def cookies
        http.cookies.select{ |c| c.session? }
    end

    #
    # Tries to find the main session (login/ID) cookie.
    #
    # @param    [Block] block   block to be passed the cookie
    #
    # @raise    [Error::NoLoginCheck]   If no login-check has been configured.
    #
    def cookie( &block )
        return block.call( @session_cookie ) if @session_cookie
        fail Error::NoLoginCheck, 'No login-check has been configured.' if !has_login_check?

        cookies.each do |cookie|
            logged_in?( cookies: { cookie.name => '' } ) do |bool|
                next if bool
                block.call( @session_cookie = cookie )
            end
        end
    end

    #
    # Finds a login forms based on supplied location, collection and criteria.
    #
    # @param    [Hash] opts
    # @option opts [Bool] :requires_password Does the login form include a password field? (Defaults to `true`)
    # @option opts [Array, Regexp] :action Regexp to match or String to compare against the form action.
    # @option opts [String, Array, Hash, Symbol] :inputs Inputs that the form must contain.
    # @option opts [Array<Element::Form>] :forms Collection of forms to look through.
    # @option opts [Page, Array<Page>] :pages Pages to look through.
    # @option opts [String] :url URL to fetch and look for forms.
    #
    # @param    [Block] block   if a block and a :url are given, the request
    #                               will run async and the block will be called
    #                               with the result of this method.
    #
    def find_login_form( opts = {}, &block )
        async = block_given?

        requires_password = (opts[:requires_password].nil? ? true : opts[:requires_password])

        find = proc do |cforms|
            cforms.select do |f|
                next if requires_password && !f.requires_password?

                oks   = []

                if action = opts[:action]
                    oks << !!(action.is_a?( Regexp ) ? f.action =~ action : f.action == action)
                end

                if inputs = opts[:inputs]
                    oks << f.has_inputs?( inputs )
                end

                oks.count( true ) == oks.size
            end.first
        end

        forms = if opts[:pages]
                    [opts[:pages]].flatten.map { |p| p.forms }.flatten
                elsif opts[:forms]
                    opts[:forms]
                elsif url = opts[:url]
                    http_opts = { http: { update_cookies: true } }

                    if async
                        page_from_url( url, http_opts ) { |p| block.call find.call( p.forms ) }
                    else
                        page_from_url( url, http_opts ).forms
                    end
                end

        find.call( forms || [] ) if !async
    end

    # @return   [Bool]  `true` if there is log-in capability, `false` otherwise.
    def can_login?
        has_login_sequence? && @login_check
    end

    # @return   [Bool, nil]
    #   `true` if logged-in, `false` otherwise, `nil` if there's no log-in
    #   capability.
    def ensure_logged_in
        return if !can_login?
        return true if logged_in?

        print_bad 'The scanner has been logged out.'
        print_info 'Trying to re-login...'

        LOGIN_TRIES.times do |i|
            break if login
            print_bad "Login attempt #{i+1} failed, retrying after " <<
                          "#{LOGIN_RETRY_WAIT} seconds..."
            sleep LOGIN_RETRY_WAIT
        end

        if !logged_in?
            print_bad 'Could not re-login.'
            false
        else
            print_ok 'Logged-in successfully.'
            true
        end
    end

    #
    # Uses the block in {#login_sequence} to login to the webapp.
    #
    # @return   [Bool, nil]
    #   `true` if login was successful, `false` if not, `nil` if no
    #   {#login_sequence} has been set.
    #
    def login
        login_sequence.call if has_login_sequence?
    end

    # @return   [Bool]  `true` if a login sequence exists, `false` otherwise.
    def has_login_sequence?
        !!login_sequence
    end

    #
    # Uses the block in {#login_check} to check in we're logged in to the webapp.
    #
    # @param    [Hash]   http_opts   Extra HTTP options to use for the check.
    # @param    [Block]  block
    #   If a block has been provided the check will be async and the result will
    #   be passed to it, otherwise the method will return the result.
    #
    #
    # @return   [Bool, nil]
    #   `true` if we're logged-in, `false` if not, `nil` if no
    #   {#login_sequence} has been set.
    #
    def logged_in?( http_opts = {}, &block )
        login_check.call( http_opts, block ) if has_login_check?
    end

    # @return   [Bool]  `true` if a login check exists, `false` otherwise.
    def has_login_check?
        !!login_check
    end

    def login_check( &block )
        return @login_check = block if block_given?

        if @opts.login_check_url && @opts.login_check_pattern
            set_login_check( @opts.login_check_url, @opts.login_check_pattern )
        end

        @login_check
    end

    #
    # A block used to login to the webapp.
    #
    # The block should log the framework into the webapp and return `true` on
    # success, `false` on failure.
    #
    # @param    [Block] block
    #   if a block has been given it will be set as the login sequence.
    #
    # @return   [Block]
    #
    def login_sequence( &block )
        if @login_form && !block_given?
            @login_sequence = proc do
                if !(refreshed = @login_form.refresh( update_cookies: true ))
                    print_bad 'Login form has disappeared, cannot login.'
                    next
                end

                refreshed.submit(
                    async:           false,
                    update_cookies:  true,
                    follow_location: false
                ).response
            end
        end

        return @login_sequence if !block_given?
        @login_sequence = block
    end

    #
    # Sets a login check using the provided `url` and `regexp`.
    #
    # @param    [String, #to_s]  url        URL to request.
    # @param    [String, Regexp] pattern
    #   Pattern to match against the body of the response.
    #
    def set_login_check( url, pattern )
        login_check do |opts, block|
            bool = nil
            http.get( url.to_s, opts.merge( async: !!block ) ) do |res|
                bool = !!res.body.match( pattern )
                block.call( bool ) if block
            end

            bool
        end
    end

    # @return   [HTTP]  http interface
    def http
        HTTP
    end

end
end