module csrf

import crypto.hmac
import crypto.sha256
import encoding.base64
import net.http
import net.urllib
import rand
import time
import x.vweb

@[params]
pub struct CsrfConfig {
pub:
	secret string
	// how long the random part of the csrf-token should be
	nonce_length int = 64
	// HTTP "safe" methods meaning they shouldn't alter state.
	// If a request with any of these methods is made, `protect` will always return true
	// https://datatracker.ietf.org/doc/html/rfc7231#section-4.2.1
	safe_methods []http.Method = [.get, .head, .options]
	// which hosts are allowed, enforced by checking the Origin and Referer header
	// if allowed_hosts contains '*' the check will be skipped.
	// Subdomains need to be included separately: a request from `"sub.example.com"`
	//  will be rejected when `allowed_host = ['example.com']`.
	allowed_hosts []string
	// if set to true both the Referer and Origin headers must match `allowed_hosts`
	// else if either one is valid the request is accepted
	check_origin_and_referer bool = true
	// the name of the csrf-token in the hidden html input
	token_name string = 'csrftoken'
	// the name of the cookie that contains the session id
	session_cookie string
	// cookie options
	cookie_name string        = 'csrftoken'
	same_site   http.SameSite = .same_site_strict_mode
	cookie_path string        = '/'
	// how long the cookie stays valid in seconds. Default is 30 days
	max_age       int = 60 * 60 * 24 * 30
	cookie_domain string
	// whether the cookie can be send only over HTTPS
	secure bool
}

pub struct CsrfContext {
pub mut:
	config CsrfConfig
	exempt bool
	// the csrftoken that should be placed in an html form
	csrf_token string
}

// set_token generates a new csrf_token and adds a Cookie to the response
pub fn (mut ctx CsrfContext) set_csrf_token[T](mut user_context T) string {
	ctx.csrf_token = set_token(mut user_context, ctx.config)
	return ctx.csrf_token
}

// clear the csrf token and cookie header from the context
pub fn (ctx &CsrfContext) clear_csrf_token[T](mut user_context T) {
	user_context.set_cookie(http.Cookie{
		name: config.cookie_name
		value: ''
		max_age: 0
	})
}

// csrf_token_input returns an HTML hidden input containing the csrf token
pub fn (ctx &CsrfContext) csrf_token_input() vweb.RawHtml {
	return '<input type="hidden" name="${ctx.config.token_name}" value="${ctx.csrf_token}">'
}

// middleware returns a handler that you can use with vweb's middleware
pub fn middleware[T](config CsrfConfig) vweb.MiddlewareOptions[T] {
	return vweb.MiddlewareOptions[T]{
		after: false
		handler: fn [config] [T](mut ctx T) bool {
			ctx.config = config
			if ctx.exempt {
				return true
			} else if ctx.req.method in config.safe_methods {
				return true
			} else {
				return protect(mut ctx, config)
			}
		}
	}
}

// set_token returns the csrftoken and sets an encrypted cookie with the hmac of
// `config.get_secret` and the csrftoken
pub fn set_token(mut ctx vweb.Context, config &CsrfConfig) string {
	expire_time := time.now().add_seconds(config.max_age)
	session_id := ctx.get_cookie(config.session_cookie) or { '' }

	token := generate_token(expire_time.unix_time(), session_id, config.nonce_length)
	cookie := generate_cookie(expire_time.unix_time(), token, config.secret)

	// the hmac key is set as a cookie and later validated with `app.token` that must
	// be in an html form
	ctx.set_cookie(http.Cookie{
		name: config.cookie_name
		value: cookie
		same_site: config.same_site
		http_only: true
		secure: config.secure
		path: config.cookie_path
		expires: expire_time
		max_age: config.max_age
	})

	return token
}

// protect returns false and sends an http 401 response when the csrf verification
// fails. protect will always return true if the current request method is in
// `config.safe_methods`.
pub fn protect(mut ctx vweb.Context, config &CsrfConfig) bool {
	// if the request method is a "safe" method we allow the request
	if ctx.req.method in config.safe_methods {
		return true
	}

	// check origin and referer header
	if check_origin_and_referer(ctx, config) == false {
		request_is_invalid(mut ctx)
		return false
	}

	// use the session id from the cookie, not from the csrftoken
	session_id := ctx.get_cookie(config.session_cookie) or { '' }

	actual_token := ctx.form[config.token_name] or {
		request_is_invalid(mut ctx)
		return false
	}
	// retrieve timestamp and nonce from csrftoken
	data := base64.url_decode_str(actual_token).split('.')
	println(data)
	if data.len < 3 {
		request_is_invalid(mut ctx)
		return false
	}

	// check the timestamp from the csrftoken against the current time
	// if an attacker would change the timestamp on the cookie, the token or both the
	// hmac would also change.
	now := time.now().unix_time()
	expire_timestamp := data[0].i64()
	if expire_timestamp < now {
		// token has expired
		request_is_invalid(mut ctx)
		return false
	}
	nonce := data.last()
	expected_token := base64.url_encode_str('${expire_timestamp}.${session_id}.${nonce}')

	mut actual_hash := ctx.get_cookie(config.cookie_name) or {
		request_is_invalid(mut ctx)
		return false
	}
	// old_expire := actual_hash.all_before('.')
	// actual_hash = actual_hash.replace_once (old_expire, expire_timestamp.str())

	// generate new hmac based on information in the http request
	expected_hash := generate_cookie(expire_timestamp, expected_token, config.secret)
	eprintln(actual_hash)
	eprintln(expected_hash)

	// if the new hmac matches the cookie value the request is legit
	if actual_hash != expected_hash {
		request_is_invalid(mut ctx)
		return false
	}
	eprintln('matching')

	return true
}

// check_origin_and_referer validates the `Origin` and `Referer` headers.
fn check_origin_and_referer(ctx vweb.Context, config &CsrfConfig) bool {
	// wildcard allow all hosts NOT SAFE!
	if '*' in config.allowed_hosts {
		return true
	}

	// only match host and match the full domain name
	// because lets say `allowed_host` = `['example.com']`.
	// Attackers shouldn't be able to bypass this check with the domain `example.com.attacker.com`

	origin := ctx.get_header(.origin) or { return false }
	origin_url := urllib.parse(origin) or { urllib.URL{} }

	valid_origin := origin_url.hostname() in config.allowed_hosts

	referer := ctx.get_header(.referer) or { return false }
	referer_url := urllib.parse(referer) or { urllib.URL{} }

	valid_referer := referer_url.hostname() in config.allowed_hosts

	if config.check_origin_and_referer {
		return valid_origin && valid_referer
	} else {
		return valid_origin || valid_referer
	}
}

// request_is_invalid sends an http 403 response
fn request_is_invalid(mut ctx vweb.Context) {
	ctx.res.set_status(.forbidden)
	ctx.text('Forbidden: Invalid or missing CSRF token')
}

fn generate_token(expire_time i64, session_id string, nonce_length int) string {
	nonce := rand.string_from_set('0123456789ABCDEFGHIJKLMNOPQRSTUVXYZabcdefghijklmnopqrstuvwxyz',
		nonce_length)
	token := '${expire_time}.${session_id}.${nonce}'

	return base64.url_encode_str(token)
}

// generate_cookie converts secret key based on the request context and a random
// token into an hmac key
fn generate_cookie(expire_time i64, token string, secret string) string {
	hash := base64.url_encode(hmac.new(secret.bytes(), token.bytes(), sha256.sum, sha256.block_size))
	cookie := '${expire_time}.${hash}'

	return cookie
}
