package main

import (
    "crypto/aes"
    "crypto/cipher"
    "database/sql"
    "encoding/gob"
    "fmt"
    "net"
    "sync"
    "time"

    _ "github.com/mattn/go-sqlite3"
)

var key = []byte("2f5680a7fb57ce83e8e83dbb1114ee31") // 32 bytes

var nonceStore = struct {
    sync.Mutex
    data map[string]int64
}{data: make(map[string]int64)}

const nonceExpiry = 30 * time.Second

type Request struct {
    Operation string
    Site      string
    Password  string
    Timestamp int64  // Unix time (seconds)
    Nonce     string // Random client-supplied nonce
}

type Response struct {
    Message string
}

var ipAttempts = struct {
    sync.Mutex
    data map[string][]int64
}{data: make(map[string][]int64)}

const maxAttemptsPerMinute = 5

func isRateLimited(ip string) bool {
    now := time.Now().Unix()
    ipAttempts.Lock()
    defer ipAttempts.Unlock()

    attempts := ipAttempts.data[ip]
    // Keep only the last 60 seconds of attempts
    recent := make([]int64, 0)
    for _, t := range attempts {
        if now-t <= 60 {
            recent = append(recent, t)
        }
    }
    ipAttempts.data[ip] = append(recent, now)
    return len(recent) >= maxAttemptsPerMinute
}

func isValidNonce(nonce string, timestamp int64) bool {
    now := time.Now().Unix()
    if now-timestamp > int64(nonceExpiry.Seconds()) {
        return false // Too old
    }

    nonceStore.Lock()
    defer nonceStore.Unlock()

    if t, exists := nonceStore.data[nonce]; exists {
        if now-t < int64(nonceExpiry.Seconds()) {
            return false // Replay detected
        }
    }

    nonceStore.data[nonce] = now
    return true
}

func encrypt(text string) ([]byte, error) {
    block, err := aes.NewCipher(key)
    if err != nil {
        return nil, err
    }
    gcm, err := cipher.NewGCM(block)
    if err != nil {
        return nil, err
    }
    nonce := make([]byte, gcm.NonceSize())
    return gcm.Seal(nonce, nonce, []byte(text), nil), nil
}

func decrypt(data []byte) (string, error) {
    block, err := aes.NewCipher(key)
    if err != nil {
        return "", err
    }
    gcm, err := cipher.NewGCM(block)
    if err != nil {
        return "", err
    }
    nonceSize := gcm.NonceSize()
    nonce, ciphertext := data[:nonceSize], data[nonceSize:]
    plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
    if err != nil {
        return "", err
    }
    return string(plaintext), nil
}

func handleConnection(conn net.Conn, db *sql.DB) {
    defer conn.Close()
    dec := gob.NewDecoder(conn)
    enc := gob.NewEncoder(conn)

    remoteAddr := conn.RemoteAddr().String()

    var req Request
    if err := dec.Decode(&req); err != nil {
        return
    }

    if isRateLimited(remoteAddr) {
        enc.Encode(Response{"Too many requests. Try later."})
        return
    }

    if !isValidNonce(req.Nonce, req.Timestamp) {
        enc.Encode(Response{"Invalid or replayed request"})
        return
    }

    switch req.Operation {
    case "store":
        encrypted, err := encrypt(req.Password)
        if err != nil {
            enc.Encode(Response{"Encryption failed"})
            return
        }
        _, err = db.Exec("INSERT OR REPLACE INTO accounts (site, password) VALUES (?, ?)", req.Site, encrypted)
        if err != nil {
            enc.Encode(Response{"Database error"})
            return
        }
        enc.Encode(Response{"Password stored successfully"})

    case "get":
        var encrypted []byte
        err := db.QueryRow("SELECT password FROM accounts WHERE site = ?", req.Site).Scan(&encrypted)
        if err == sql.ErrNoRows {
            enc.Encode(Response{"Site not found"})
            return
        } else if err != nil {
            enc.Encode(Response{"Database error"})
            return
        }

        pwd, err := decrypt(encrypted)
        if err != nil {
            enc.Encode(Response{"Decryption failed"})
            return
        }
        enc.Encode(Response{pwd})

    default:
        enc.Encode(Response{"Unknown operation"})
    }
}

func main() {
    db, err := sql.Open("sqlite3", "vault.db")
    if err != nil {
        panic(err)
    }
    defer db.Close()

    _, err = db.Exec(`CREATE TABLE IF NOT EXISTS accounts (
        site TEXT PRIMARY KEY,
        password BLOB
    )`)
    if err != nil {
        panic(err)
    }

    ln, err := net.Listen("tcp", ":9898")
    if err != nil {
        panic(err)
    }
    fmt.Println("Server started on port 9898")

    for {
        conn, err := ln.Accept()
        if err != nil {
            continue
        }
        go handleConnection(conn, db)
    }
}