// Copyright (c) 2019-2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by a GPL license that can be found in the LICENSE file.
module main

import veb
import time
import crypto.hmac
import crypto.sha256
import encoding.hex

const two_factor_pending_cookie = 'pending_2fa'
const two_factor_pending_ttl = 300 // seconds

fn (mut app App) pending_2fa_key(user User) []u8 {
	return '${user.password}:${user.salt}'.bytes()
}

fn (mut app App) sign_pending_2fa(user User, expires i64) string {
	payload := '${user.id}:${expires}'
	mac := hmac.new(app.pending_2fa_key(user), payload.bytes(), sha256.sum, sha256.block_size)
	return '${payload}:${hex.encode(mac)}'
}

fn (mut app App) verify_pending_2fa(token string) ?User {
	parts := token.split(':')
	if parts.len != 3 {
		return none
	}
	user_id := parts[0].int()
	expires := parts[1].i64()
	sig := parts[2]
	if expires < time.now().unix() {
		return none
	}
	user := app.get_user_by_id(user_id) or { return none }
	payload := '${user_id}:${expires}'
	expected := hmac.new(app.pending_2fa_key(user), payload.bytes(), sha256.sum, sha256.block_size)
	if hex.encode(expected) != sig {
		return none
	}
	return user
}

@['/login/2fa']
pub fn (mut app App) two_factor_prompt(mut ctx Context) veb.Result {
	pending := ctx.get_cookie(two_factor_pending_cookie) or { return ctx.redirect_to_login() }
	app.verify_pending_2fa(pending) or { return ctx.redirect_to_login() }
	return $veb.html('templates/two_factor_login.html')
}

@['/login/2fa'; post]
pub fn (mut app App) handle_two_factor_login(mut ctx Context, code string) veb.Result {
	pending := ctx.get_cookie(two_factor_pending_cookie) or { return ctx.redirect_to_login() }
	user := app.verify_pending_2fa(pending) or { return ctx.redirect_to_login() }
	tf := app.find_two_factor(user.id) or { return ctx.redirect_to_login() }
	if !tf.is_enabled || !verify_totp(tf.secret, code.trim_space()) {
		ctx.error('Invalid verification code')
		return $veb.html('templates/two_factor_login.html')
	}
	ctx.set_cookie(name: two_factor_pending_cookie, value: '')
	app.auth_user(mut ctx, user, ctx.ip()) or {
		ctx.error('There was an error while logging in')
		return ctx.redirect_to_login()
	}
	app.add_security_log(user_id: user.id, kind: .logged_in) or { app.info(err.str()) }
	return ctx.redirect('/${user.username}')
}

@['/:username/settings/2fa']
pub fn (mut app App) view_two_factor_settings(mut ctx Context, username string) veb.Result {
	if !ctx.logged_in || ctx.user.username != username {
		return ctx.redirect_to_index()
	}
	tf := app.find_two_factor(ctx.user.id) or {
		TwoFactor{
			user_id: ctx.user.id
		}
	}
	enabled := tf.is_enabled
	mut secret := ''
	mut provisioning_uri := ''
	if !enabled {
		secret = if tf.secret == '' { generate_totp_secret() } else { tf.secret }
		app.upsert_two_factor(ctx.user.id, secret, false) or {}
		provisioning_uri = totp_provisioning_uri(username, secret)
	}
	return $veb.html('templates/two_factor_settings.html')
}

@['/:username/settings/2fa/enable'; post]
pub fn (mut app App) handle_enable_two_factor(mut ctx Context, username string) veb.Result {
	if !ctx.logged_in || ctx.user.username != username {
		return ctx.redirect_to_index()
	}
	code := ctx.form['code'].trim_space()
	tf := app.find_two_factor(ctx.user.id) or { return ctx.redirect('/${username}/settings/2fa') }
	if !verify_totp(tf.secret, code) {
		ctx.error('Invalid verification code')
		return ctx.redirect('/${username}/settings/2fa')
	}
	app.upsert_two_factor(ctx.user.id, tf.secret, true) or {}
	return ctx.redirect('/${username}/settings/2fa')
}

@['/:username/settings/2fa/disable'; post]
pub fn (mut app App) handle_disable_two_factor(mut ctx Context, username string) veb.Result {
	if !ctx.logged_in || ctx.user.username != username {
		return ctx.redirect_to_index()
	}
	app.delete_two_factor(ctx.user.id) or {}
	return ctx.redirect('/${username}/settings/2fa')
}