package main

import (
    "golang.org/x/crypto/bcrypt"
    "golang.org/x/crypto/chacha20poly1305"
    "golang.org/x/crypto/scrypt"
    "crypto/sha256"
    "crypto/rand"
    "crypto/aes"
    "crypto/cipher"
    "database/sql"
    "encoding/base64"
    "encoding/pem"
    "encoding/gob"
    "fmt"
    "net"
    "sync"
    "time"
    "errors"
    "flag"
    "os"
    "log"
    "io/ioutil"
    "gopkg.in/yaml.v2"
    
    _ "github.com/mattn/go-sqlite3"
)

var ChaKey = []byte("")

type Config struct {
    AllowRegistration bool `yaml:"AllowRegistration"`
	Auth struct {
		ChaKey string `yaml:"ChaKey"`
        PEM    string `yaml:"PEM"`
	} `yaml:"auth"`
}

var config Config

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

const nonceExpiry = 30 * time.Second

var key []byte
type Request struct {
    KeyPwd    string
    Username  string
    Password  string
    Operation string
    Site      string
    VaultData string
    Timestamp int64  // Unix time (seconds)
    Nonce     string // Random client-supplied nonce
}

type Response struct {
    Message string
    Enc     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
}

// decryptPEMToKey decrypts a PEM block with the given password
func decryptPEMToKey(filename string, password string) ([]byte, error) {
    // Read the entire file
    pemData, err := os.ReadFile(filename)
    if err != nil {
        return nil, fmt.Errorf("failed to read PEM file: %v", err)
    }
    
    // Decode the PEM block
    pemBlock, _ := pem.Decode(pemData)
    if pemBlock == nil {
        return nil, errors.New("failed to decode PEM block")
    }    
    
	// Check if the PEM block is encrypted
	if pemBlock.Headers["Key-Derivation"] == "" {
		return nil, errors.New("PEM block is not password-protected")
	}

	// Get salt and IV from headers
	salt, err := hexStringToBytes(pemBlock.Headers["SALT"])
	if err != nil {
		return nil, fmt.Errorf("invalid salt: %v", err)
	}

	iv, err := hexStringToBytes(pemBlock.Headers["IV"])
	if err != nil {
		return nil, fmt.Errorf("invalid IV: %v", err)
	}

	// Derive encryption key
	var encryptionKey []byte
	switch pemBlock.Headers["Key-Derivation"] {
	case "scrypt":
		encryptionKey, err = scrypt.Key([]byte(password), salt, 32768, 8, 1, 32)
	//case "pbkdf2":
		//encryptionKey = pbkdf2.Key([]byte(password), salt, 10000, 32, sha256.New)
	default:
		return nil, errors.New("unsupported key derivation function")
	}
	if err != nil {
		return nil, err
	}

	// Decrypt the key
	block, err := aes.NewCipher(encryptionKey)
	if err != nil {
		return nil, err
	}

	decrypted := make([]byte, len(pemBlock.Bytes))
	stream := cipher.NewCFBDecrypter(block, iv)
	stream.XORKeyStream(decrypted, pemBlock.Bytes)

	return decrypted, nil
}

// hexStringToBytes converts a hex string to bytes
func hexStringToBytes(hexStr string) ([]byte, error) {
	if len(hexStr)%2 != 0 {
		return nil, errors.New("hex string length must be even")
	}
	bytes := make([]byte, len(hexStr)/2)
	for i := 0; i < len(hexStr); i += 2 {
		b, err := hexByte(hexStr[i], hexStr[i+1])
		if err != nil {
			return nil, err
		}
		bytes[i/2] = b
	}
	return bytes, nil
}

// hexByte converts two hex characters to a byte
func hexByte(a, b byte) (byte, error) {
	high, err := hexDigitToByte(a)
	if err != nil {
		return 0, err
	}
	low, err := hexDigitToByte(b)
	if err != nil {
		return 0, err
	}
	return (high << 4) | low, nil
}

// hexDigitToByte converts a single hex character to its byte value
func hexDigitToByte(c byte) (byte, error) {
	switch {
	case '0' <= c && c <= '9':
		return c - '0', nil
	case 'a' <= c && c <= 'f':
		return c - 'a' + 10, nil
	case 'A' <= c && c <= 'F':
		return c - 'A' + 10, nil
	default:
		return 0, fmt.Errorf("invalid hex digit %c", c)
	}
}

func hashPassword(pw string) (string, error) {
    bytes, err := bcrypt.GenerateFromPassword([]byte(pw), 12)
    return string(bytes), err
}

func checkPassword(hash, pw string) error {
    return bcrypt.CompareHashAndPassword([]byte(hash), []byte(pw))
}

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

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

func getKey()([]byte) {
    // Hash it to 32 bytes using SHA-256
    hashedKey := sha256.Sum256(ChaKey)
    theKey := hashedKey[:] // Convert [32]byte to []byte
    return theKey
}

func chEnc(pwd string)(string, error) {
    // Create cipher
    aead, err := chacha20poly1305.NewX(getKey())
    if err != nil {
        return "", err
    }    
    // Create nonce
    nonce := make([]byte, aead.NonceSize())
    if _, err := rand.Read(nonce); err != nil {
        return "", err
    }    
        
    // Encrypt the encoded data
    encrypted := aead.Seal(nil, nonce, []byte(pwd), nil)
        
    // Send nonce + encrypted data
    fullMessage := append(nonce, encrypted...)
    encoded := base64.StdEncoding.EncodeToString(fullMessage)

    return encoded, nil
}

func chDec(eText string)(string) {
    if eText == "" {
       log.Fatalf("Error: Blank")
    }
    
    // Decoding from base64
	encryptedMsg, err := base64.StdEncoding.DecodeString(eText)
	if err != nil {
		log.Fatalf("Error: Base64 decode")
	}
    
	// Create cipher instance (XChaCha20-Poly1305 for longer nonces)
	aead, err := chacha20poly1305.NewX(getKey())
	if err != nil {
		log.Fatalf("Error: Cha20 key")
	}

	// Decrypt: Split nonce and ciphertext
	decryptedNonce := encryptedMsg[:aead.NonceSize()]
	decryptedCiphertext := encryptedMsg[aead.NonceSize():]

	decrypted, err := aead.Open(nil, decryptedNonce, decryptedCiphertext, nil)
	if err != nil {
		log.Fatalf("Error: aead open")
	}
    return string(decrypted)
}

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
    }
    
    username := chDec(req.Username)
    keypwd := chDec(req.KeyPwd)
    password := chDec(req.Password)
    
    if username == "" {
        enc.Encode(Response{Message: "Required username, blank found!", Enc: ""})
        return
    }
    if keypwd == "" {
        enc.Encode(Response{Message: "Required key passphrase, blank found!", Enc: ""})
        return
    }
    if password == "" {
        enc.Encode(Response{Message: "Required user password, blank found!", Enc: ""})
        return
    }
    
    k, err := decryptPEMToKey(config.Auth.PEM, keypwd)
    if err != nil {
        enc.Encode(Response{Message: "Unable to decode data!", Enc: ""})
        return
    }
    key = k

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

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

    if req.Operation == "register" {
        if config.AllowRegistration == false {
            enc.Encode(Response{Message: "Registration Disabled!", Enc: ""})
            return
        }
        hashed, err := hashPassword(password)
        if err != nil {
            enc.Encode(Response{Message: "Failed to hash password", Enc: ""})
            return
        }
        _, err = db.Exec("INSERT INTO users (username, password) VALUES (?, ?)", username, hashed)
        if err != nil {
            enc.Encode(Response{Message: "Registration failed (user exists?)", Enc: ""})
        } else {
            enc.Encode(Response{Message: "Registration successful", Enc: ""})
        }
        return
    }
    
    site := chDec(req.Site)
    if site == "" {
        enc.Encode(Response{Message: "Required site name, blank found!", Enc: ""})
        return
    }
    
    // Authenticate all other operations
    var storedHash string
    err = db.QueryRow("SELECT password FROM users WHERE username = ?", username).Scan(&storedHash)
    if err != nil || checkPassword(storedHash, password) != nil {
        enc.Encode(Response{Message: "Authentication failed", Enc: ""})
        return
    }

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

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

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

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

func main() {
    vaultPtr := flag.String("dbfile", "vault.db", "Enter dbfile")
    configFilePtr := flag.String("config", "config.yaml", "Path to the YAML configuration file")
	flag.Parse()
    
    if *vaultPtr == "" {
		fmt.Println("Please specify a db file using -dbfile")
		flag.Usage()
		os.Exit(1)
	}
    
    if *configFilePtr == "" {
		fmt.Println("Please specify a configuration file using -config")
		flag.Usage()
		os.Exit(1)
	}

	// Read the config file
	yamlFile, err := ioutil.ReadFile(*configFilePtr)
	if err != nil {
		log.Fatalf("Error reading YAML file: %v\n", err)
	}

	// Parse the YAML
	err = yaml.Unmarshal(yamlFile, &config)
	if err != nil {
		log.Fatalf("Error parsing YAML file: %v\n", err)
	}

	// Print the ChaKey
	if config.Auth.ChaKey == "" {
		fmt.Println("Warning: ChaKey not found in configuration file")
    }
    ChaKey = []byte(config.Auth.ChaKey)
    if config.Auth.PEM == "" {
        fmt.Println("Warning: PEM file not found in configuration file")
    }
            
    db, err := sql.Open("sqlite3", *vaultPtr)
    if err != nil {
        panic(err)
    }
    defer db.Close()

    _, err = db.Exec(`
CREATE TABLE IF NOT EXISTS users (
    username TEXT PRIMARY KEY,
    password TEXT
);

CREATE TABLE IF NOT EXISTS accounts (
    user TEXT,
    site TEXT,
    password BLOB,
    PRIMARY KEY(user, site),
    FOREIGN KEY(user) REFERENCES users(username)
)
`)
    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)
    }
}