414 lines
8.9 KiB
Go
414 lines
8.9 KiB
Go
package main
|
|
|
|
import (
|
|
"database/sql"
|
|
"flag"
|
|
"fmt"
|
|
"html/template"
|
|
"log"
|
|
"net/http"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/golang-jwt/jwt/v5"
|
|
"github.com/pquerna/otp/totp"
|
|
|
|
_ "github.com/mattn/go-sqlite3"
|
|
)
|
|
|
|
const (
|
|
dbFile = "auth.db"
|
|
jwtCookie = "auth_token"
|
|
sessionDuration = 24 * time.Hour
|
|
maxLoginAttempts = 5
|
|
rateLimitWindow = 15 * time.Minute
|
|
)
|
|
|
|
var (
|
|
db *sql.DB
|
|
jwtSecret string
|
|
|
|
// Rate limiting
|
|
loginAttempts = make(map[string]*attemptTracker)
|
|
attemptsMux sync.RWMutex
|
|
)
|
|
|
|
type attemptTracker struct {
|
|
count int
|
|
firstAttempt time.Time
|
|
}
|
|
|
|
type Claims struct {
|
|
jwt.RegisteredClaims
|
|
}
|
|
|
|
func main() {
|
|
generate := flag.Bool("generate", false, "Generate a new TOTP seed and add it to the database")
|
|
flag.Parse()
|
|
|
|
// Load JWT secret from environment
|
|
jwtSecret = os.Getenv("JWT_SECRET")
|
|
if jwtSecret == "" {
|
|
log.Fatal("JWT_SECRET environment variable must be set!")
|
|
}
|
|
|
|
var err error
|
|
db, err = sql.Open("sqlite3", dbFile)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
defer db.Close()
|
|
|
|
_, err = db.Exec(`CREATE TABLE IF NOT EXISTS seeds (id INTEGER PRIMARY KEY, secret TEXT UNIQUE)`)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
if *generate {
|
|
generateSeed()
|
|
os.Exit(0)
|
|
}
|
|
|
|
// Check if there are any seeds, if not, generate one
|
|
var count int
|
|
err = db.QueryRow(`SELECT COUNT(*) FROM seeds`).Scan(&count)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
if count == 0 {
|
|
log.Println("No seeds found, generating one...")
|
|
generateSeed()
|
|
}
|
|
|
|
// Start cleanup goroutine for rate limiting
|
|
go cleanupRateLimits()
|
|
|
|
http.HandleFunc("/verify", verifyHandler)
|
|
http.HandleFunc("/login", loginHandler)
|
|
http.HandleFunc("/health", healthHandler)
|
|
|
|
log.Println("Starting auth server on :3000")
|
|
log.Fatal(http.ListenAndServe(":3000", nil))
|
|
}
|
|
|
|
func generateSeed() {
|
|
key, err := totp.Generate(totp.GenerateOpts{
|
|
Issuer: "ForwardAuthApp",
|
|
AccountName: "user",
|
|
})
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
secret := key.Secret()
|
|
_, err = db.Exec(`INSERT INTO seeds (secret) VALUES (?)`, secret)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
fmt.Printf("New TOTP seed generated:\nSecret: %s\nOTPAuth URL: %s\n", secret, key.URL())
|
|
fmt.Println("Use this to set up your authenticator app.")
|
|
}
|
|
|
|
func cleanupRateLimits() {
|
|
ticker := time.NewTicker(5 * time.Minute)
|
|
defer ticker.Stop()
|
|
for range ticker.C {
|
|
attemptsMux.Lock()
|
|
now := time.Now()
|
|
for ip, tracker := range loginAttempts {
|
|
if now.Sub(tracker.firstAttempt) > rateLimitWindow {
|
|
delete(loginAttempts, ip)
|
|
}
|
|
}
|
|
attemptsMux.Unlock()
|
|
}
|
|
}
|
|
|
|
func checkRateLimit(ip string) bool {
|
|
attemptsMux.Lock()
|
|
defer attemptsMux.Unlock()
|
|
|
|
tracker, exists := loginAttempts[ip]
|
|
if !exists {
|
|
return true
|
|
}
|
|
|
|
// Reset if window expired
|
|
if time.Since(tracker.firstAttempt) > rateLimitWindow {
|
|
delete(loginAttempts, ip)
|
|
return true
|
|
}
|
|
|
|
return tracker.count < maxLoginAttempts
|
|
}
|
|
|
|
func recordFailedAttempt(ip string) {
|
|
attemptsMux.Lock()
|
|
defer attemptsMux.Unlock()
|
|
|
|
tracker, exists := loginAttempts[ip]
|
|
if !exists {
|
|
loginAttempts[ip] = &attemptTracker{
|
|
count: 1,
|
|
firstAttempt: time.Now(),
|
|
}
|
|
return
|
|
}
|
|
|
|
// Reset if window expired
|
|
if time.Since(tracker.firstAttempt) > rateLimitWindow {
|
|
tracker.count = 1
|
|
tracker.firstAttempt = time.Now()
|
|
return
|
|
}
|
|
|
|
tracker.count++
|
|
}
|
|
|
|
func clearRateLimit(ip string) {
|
|
attemptsMux.Lock()
|
|
defer attemptsMux.Unlock()
|
|
delete(loginAttempts, ip)
|
|
}
|
|
|
|
func healthHandler(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusOK)
|
|
w.Write([]byte("OK"))
|
|
}
|
|
|
|
func verifyHandler(w http.ResponseWriter, r *http.Request) {
|
|
tokenStr, err := getCookie(r, jwtCookie)
|
|
if err == nil {
|
|
_, err = validateJWT(tokenStr)
|
|
if err == nil {
|
|
w.WriteHeader(http.StatusNoContent) // 204
|
|
return
|
|
}
|
|
}
|
|
|
|
// Not authenticated, redirect to login
|
|
next := r.Header.Get("X-Original-URI")
|
|
if next == "" {
|
|
next = "/"
|
|
}
|
|
http.Redirect(w, r, "/login?next="+next, http.StatusFound) // 302
|
|
}
|
|
|
|
var loginTmpl = template.Must(template.New("login").Parse(`
|
|
<!DOCTYPE html>
|
|
<html>
|
|
<head>
|
|
<title>Login</title>
|
|
<style>
|
|
body {
|
|
font-family: Arial, sans-serif;
|
|
display: flex;
|
|
justify-content: center;
|
|
align-items: center;
|
|
min-height: 100vh;
|
|
margin: 0;
|
|
background: #f0f0f0;
|
|
}
|
|
.login-container {
|
|
background: white;
|
|
padding: 2rem;
|
|
border-radius: 8px;
|
|
box-shadow: 0 2px 10px rgba(0,0,0,0.1);
|
|
width: 100%;
|
|
max-width: 400px;
|
|
}
|
|
h1 {
|
|
margin-top: 0;
|
|
color: #333;
|
|
}
|
|
.error {
|
|
color: #d32f2f;
|
|
background: #ffebee;
|
|
padding: 0.75rem;
|
|
border-radius: 4px;
|
|
margin-bottom: 1rem;
|
|
}
|
|
form {
|
|
display: flex;
|
|
flex-direction: column;
|
|
}
|
|
label {
|
|
margin-bottom: 0.5rem;
|
|
color: #555;
|
|
}
|
|
input[type="text"] {
|
|
padding: 0.75rem;
|
|
border: 1px solid #ddd;
|
|
border-radius: 4px;
|
|
font-size: 1rem;
|
|
margin-bottom: 1rem;
|
|
}
|
|
button {
|
|
padding: 0.75rem;
|
|
background: #1976d2;
|
|
color: white;
|
|
border: none;
|
|
border-radius: 4px;
|
|
font-size: 1rem;
|
|
cursor: pointer;
|
|
}
|
|
button:hover {
|
|
background: #1565c0;
|
|
}
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div class="login-container">
|
|
<h1>Enter OTP</h1>
|
|
{{if .Error}}
|
|
<div class="error">{{.Error}}</div>
|
|
{{end}}
|
|
<form method="post" action="/login">
|
|
<input type="hidden" name="next" value="{{.Next}}">
|
|
<label for="otp">One-Time Password:</label>
|
|
<input type="text" id="otp" name="otp" required autofocus autocomplete="off" pattern="[0-9]{6}" maxlength="6" placeholder="000000">
|
|
<button type="submit">Submit</button>
|
|
</form>
|
|
</div>
|
|
</body>
|
|
</html>
|
|
`))
|
|
|
|
type loginData struct {
|
|
Next string
|
|
Error string
|
|
}
|
|
|
|
func loginHandler(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method == "GET" {
|
|
next := r.URL.Query().Get("next")
|
|
data := loginData{Next: next}
|
|
loginTmpl.Execute(w, data)
|
|
return
|
|
}
|
|
|
|
if r.Method != "POST" {
|
|
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
|
|
// Rate limiting check
|
|
clientIP := r.RemoteAddr
|
|
if !checkRateLimit(clientIP) {
|
|
data := loginData{
|
|
Next: r.URL.Query().Get("next"),
|
|
Error: "Too many failed attempts. Try again in 15 minutes.",
|
|
}
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
loginTmpl.Execute(w, data)
|
|
return
|
|
}
|
|
|
|
err := r.ParseForm()
|
|
if err != nil {
|
|
http.Error(w, "Bad request", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
otpCode := strings.TrimSpace(r.FormValue("otp"))
|
|
next := r.FormValue("next")
|
|
if next == "" {
|
|
next = "/"
|
|
}
|
|
|
|
// Validate redirect target to prevent open redirects
|
|
if !strings.HasPrefix(next, "/") {
|
|
next = "/"
|
|
}
|
|
|
|
if validateOTP(otpCode) {
|
|
tokenStr, err := generateJWT()
|
|
if err != nil {
|
|
log.Printf("Error generating JWT: %v", err)
|
|
http.Error(w, "Internal error", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
// Clear rate limit on successful login
|
|
clearRateLimit(clientIP)
|
|
|
|
http.SetCookie(w, &http.Cookie{
|
|
Name: jwtCookie,
|
|
Value: tokenStr,
|
|
Expires: time.Now().Add(sessionDuration),
|
|
HttpOnly: true,
|
|
SameSite: http.SameSiteLaxMode,
|
|
Secure: true,
|
|
Path: "/",
|
|
})
|
|
|
|
http.Redirect(w, r, next, http.StatusFound)
|
|
return
|
|
}
|
|
|
|
// Invalid OTP - record failed attempt
|
|
recordFailedAttempt(clientIP)
|
|
data := loginData{Next: next, Error: "Invalid OTP"}
|
|
w.WriteHeader(http.StatusUnauthorized)
|
|
loginTmpl.Execute(w, data)
|
|
}
|
|
|
|
func validateOTP(otpCode string) bool {
|
|
rows, err := db.Query(`SELECT secret FROM seeds`)
|
|
if err != nil {
|
|
log.Printf("Error querying seeds: %v", err)
|
|
return false
|
|
}
|
|
defer rows.Close()
|
|
|
|
for rows.Next() {
|
|
var secret string
|
|
err = rows.Scan(&secret)
|
|
if err != nil {
|
|
log.Printf("Error scanning seed: %v", err)
|
|
continue
|
|
}
|
|
if totp.Validate(otpCode, secret) {
|
|
return true
|
|
}
|
|
}
|
|
|
|
if err = rows.Err(); err != nil {
|
|
log.Printf("Error iterating rows: %v", err)
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
func generateJWT() (string, error) {
|
|
claims := Claims{
|
|
RegisteredClaims: jwt.RegisteredClaims{
|
|
ExpiresAt: jwt.NewNumericDate(time.Now().Add(sessionDuration)),
|
|
IssuedAt: jwt.NewNumericDate(time.Now()),
|
|
},
|
|
}
|
|
token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
|
|
return token.SignedString([]byte(jwtSecret))
|
|
}
|
|
|
|
func validateJWT(tokenStr string) (*Claims, error) {
|
|
claims := &Claims{}
|
|
token, err := jwt.ParseWithClaims(tokenStr, claims, func(token *jwt.Token) (interface{}, error) {
|
|
return []byte(jwtSecret), nil
|
|
})
|
|
if err != nil || !token.Valid {
|
|
return nil, err
|
|
}
|
|
return claims, nil
|
|
}
|
|
|
|
func getCookie(r *http.Request, name string) (string, error) {
|
|
cookie, err := r.Cookie(name)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return cookie.Value, nil
|
|
} |