Wormhole/internal/store/store.go

package store

import (
	"fmt"
	"os"
	"path/filepath"
	"time"

	"github.com/google/uuid"
	"github.com/glebarez/sqlite"
	"golang.org/x/crypto/bcrypt"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

type Store struct {
	db *gorm.DB
}

func Open(path string) (*Store, error) {
	dir := filepath.Dir(path)
	if err := os.MkdirAll(dir, 0755); err != nil {
		return nil, err
	}
	db, err := gorm.Open(sqlite.Open(path), &gorm.Config{
		Logger: logger.Default.LogMode(logger.Warn),
	})
	if err != nil {
		return nil, err
	}
	sqlDB, err := db.DB()
	if err != nil {
		return nil, err
	}
	sqlDB.SetMaxOpenConns(1)
	if _, err := sqlDB.Exec("PRAGMA journal_mode=WAL"); err != nil {
		return nil, fmt.Errorf("enable WAL: %w", err)
	}
	if _, err := sqlDB.Exec("PRAGMA busy_timeout=5000"); err != nil {
		return nil, fmt.Errorf("set busy_timeout: %w", err)
	}
	s := &Store{db: db}
	if err := s.migrate(); err != nil {
		return nil, err
	}
	if err := s.seed(); err != nil {
		return nil, err
	}
	return s, nil
}

func (s *Store) DB() *gorm.DB {
	return s.db
}

func (s *Store) migrate() error {
	return s.db.AutoMigrate(
		&User{},
		&Client{},
		&Tunnel{},
		&Host{},
		&IPRule{},
		&AuthPolicy{},
		&FlowStat{},
		&RequestIP{},
		&Session{},
		&SystemSettings{},
		&APIKey{},
		&ResourceVisitor{},
		&VisitorAccessLog{},
	)
}

func (s *Store) seed() error {
	var count int64
	s.db.Model(&User{}).Count(&count)
	if count > 0 {
		return nil
	}
	hash, err := bcrypt.GenerateFromPassword([]byte("admin"), bcrypt.DefaultCost)
	if err != nil {
		return err
	}
	admin := User{
		Username:     "admin",
		PasswordHash: string(hash),
		Role:         "admin",
		Status:       true,
	}
	return s.db.Create(&admin).Error
}

func GenerateVerifyKey() string {
	return uuid.New().String()
}

func HashPassword(password string) (string, error) {
	hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return "", err
	}
	return string(hash), nil
}

func CheckPassword(hash, password string) bool {
	return bcrypt.CompareHashAndPassword([]byte(hash), []byte(password)) == nil
}

func (s *Store) ListClients(ownerID uint, isAdmin bool) ([]Client, error) {
	var clients []Client
	q := s.db.Order("id asc")
	if !isAdmin {
		q = q.Where("owner_user_id = ?", ownerID)
	}
	return clients, q.Find(&clients).Error
}

func (s *Store) GetClientByID(id uint) (*Client, error) {
	var c Client
	if err := s.db.First(&c, id).Error; err != nil {
		return nil, err
	}
	return &c, nil
}

func (s *Store) GetClientByVerifyKey(key string) (*Client, error) {
	var c Client
	if err := s.db.Where("verify_key = ? AND status = ?", key, true).First(&c).Error; err != nil {
		return nil, err
	}
	return &c, nil
}

func (s *Store) CreateClient(c *Client) error {
	if c.VerifyKey == "" {
		c.VerifyKey = GenerateVerifyKey()
	}
	return s.db.Create(c).Error
}

func (s *Store) UpdateClient(c *Client) error {
	return s.db.Save(c).Error
}

func (s *Store) DeleteClient(id uint) error {
	return s.db.Transaction(func(tx *gorm.DB) error {
		var tunnelIDs, hostIDs []uint
		tx.Model(&Tunnel{}).Where("client_id = ?", id).Pluck("id", &tunnelIDs)
		tx.Model(&Host{}).Where("client_id = ?", id).Pluck("id", &hostIDs)

		for _, tid := range tunnelIDs {
			if err := tx.Where("resource_type = ? AND resource_id = ?", "tunnel", tid).Delete(&AuthPolicy{}).Error; err != nil {
				return err
			}
			if err := tx.Where("resource_type = ? AND resource_id = ?", "tunnel", tid).Delete(&ResourceVisitor{}).Error; err != nil {
				return err
			}
			if err := tx.Where("scope = ? AND resource_id = ?", "tunnel", tid).Delete(&IPRule{}).Error; err != nil {
				return err
			}
			if err := tx.Where("resource_type = ? AND resource_id = ?", "tunnel", tid).Delete(&FlowStat{}).Error; err != nil {
				return err
			}
			if err := tx.Where("resource_type = ? AND resource_id = ?", "tunnel", tid).Delete(&RequestIP{}).Error; err != nil {
				return err
			}
		}
		for _, hid := range hostIDs {
			if err := tx.Where("resource_type = ? AND resource_id = ?", "host", hid).Delete(&AuthPolicy{}).Error; err != nil {
				return err
			}
			if err := tx.Where("resource_type = ? AND resource_id = ?", "host", hid).Delete(&ResourceVisitor{}).Error; err != nil {
				return err
			}
			if err := tx.Where("scope = ? AND resource_id = ?", "host", hid).Delete(&IPRule{}).Error; err != nil {
				return err
			}
			if err := tx.Where("resource_type = ? AND resource_id = ?", "host", hid).Delete(&FlowStat{}).Error; err != nil {
				return err
			}
			if err := tx.Where("resource_type = ? AND resource_id = ?", "host", hid).Delete(&RequestIP{}).Error; err != nil {
				return err
			}
		}
		if err := tx.Where("scope = ? AND resource_id = ?", "client", id).Delete(&IPRule{}).Error; err != nil {
			return err
		}
		if err := tx.Where("client_id = ?", id).Delete(&Tunnel{}).Error; err != nil {
			return err
		}
		if err := tx.Where("client_id = ?", id).Delete(&Host{}).Error; err != nil {
			return err
		}
		return tx.Delete(&Client{}, id).Error
	})
}

func (s *Store) ListTunnels(clientID uint, ownerID uint, isAdmin bool) ([]Tunnel, error) {
	var tunnels []Tunnel
	q := s.db.Order("id asc")
	if clientID > 0 {
		q = q.Where("client_id = ?", clientID)
	} else if !isAdmin {
		var ids []uint
		s.db.Model(&Client{}).Where("owner_user_id = ?", ownerID).Pluck("id", &ids)
		if len(ids) == 0 {
			return tunnels, nil
		}
		q = q.Where("client_id IN ?", ids)
	}
	return tunnels, q.Find(&tunnels).Error
}

func (s *Store) GetTunnelByID(id uint) (*Tunnel, error) {
	var t Tunnel
	if err := s.db.First(&t, id).Error; err != nil {
		return nil, err
	}
	return &t, nil
}

func (s *Store) CreateTunnel(t *Tunnel) error {
	return s.db.Create(t).Error
}

func (s *Store) UpdateTunnel(t *Tunnel) error {
	return s.db.Save(t).Error
}

func (s *Store) DeleteTunnel(id uint) error {
	_ = s.DeleteResourceVisitors("tunnel", id)
	return s.db.Delete(&Tunnel{}, id).Error
}

func (s *Store) ListHosts(clientID uint, ownerID uint, isAdmin bool) ([]Host, error) {
	var hosts []Host
	q := s.db.Order("id asc")
	if clientID > 0 {
		q = q.Where("client_id = ?", clientID)
	} else if !isAdmin {
		var ids []uint
		s.db.Model(&Client{}).Where("owner_user_id = ?", ownerID).Pluck("id", &ids)
		if len(ids) == 0 {
			return hosts, nil
		}
		q = q.Where("client_id IN ?", ids)
	}
	return hosts, q.Find(&hosts).Error
}

func (s *Store) GetHostByID(id uint) (*Host, error) {
	var h Host
	if err := s.db.First(&h, id).Error; err != nil {
		return nil, err
	}
	return &h, nil
}

func (s *Store) CreateHost(h *Host) error {
	return s.db.Create(h).Error
}

func (s *Store) UpdateHost(h *Host) error {
	return s.db.Save(h).Error
}

func (s *Store) DeleteHost(id uint) error {
	_ = s.DeleteResourceVisitors("host", id)
	return s.db.Delete(&Host{}, id).Error
}

func (s *Store) ListIPRules(scope string, resourceID uint) ([]IPRule, error) {
	var rules []IPRule
	q := s.db.Order("priority desc, id asc")
	if scope != "" {
		q = q.Where("scope = ?", scope)
	}
	if resourceID > 0 {
		q = q.Where("resource_id = ?", resourceID)
	}
	return rules, q.Find(&rules).Error
}

func (s *Store) CreateIPRule(r *IPRule) error {
	return s.db.Create(r).Error
}

func (s *Store) DeleteIPRule(id uint) error {
	return s.db.Delete(&IPRule{}, id).Error
}

func (s *Store) ListAuthPolicies(resourceType string, resourceID uint) ([]AuthPolicy, error) {
	var policies []AuthPolicy
	q := s.db.Order("id desc")
	if resourceType != "" {
		q = q.Where("resource_type = ?", resourceType)
	}
	if resourceID > 0 {
		q = q.Where("resource_id = ?", resourceID)
	}
	return policies, q.Find(&policies).Error
}

func (s *Store) GetAuthPolicy(id uint) (*AuthPolicy, error) {
	var p AuthPolicy
	if err := s.db.First(&p, id).Error; err != nil {
		return nil, err
	}
	return &p, nil
}

func (s *Store) CreateAuthPolicy(p *AuthPolicy) error {
	return s.db.Create(p).Error
}

func (s *Store) UpdateAuthPolicy(p *AuthPolicy) error {
	return s.db.Save(p).Error
}

func (s *Store) DeleteAuthPolicy(id uint) error {
	return s.db.Delete(&AuthPolicy{}, id).Error
}

func (s *Store) ListUsers() ([]User, error) {
	var users []User
	return users, s.db.Order("id asc").Find(&users).Error
}

func (s *Store) GetUserByUsername(username string) (*User, error) {
	var u User
	if err := s.db.Where("username = ?", username).First(&u).Error; err != nil {
		return nil, err
	}
	return &u, nil
}

func (s *Store) GetUserByID(id uint) (*User, error) {
	var u User
	if err := s.db.First(&u, id).Error; err != nil {
		return nil, err
	}
	return &u, nil
}

func (s *Store) CreateUser(u *User) error {
	return s.db.Create(u).Error
}

func (s *Store) UpdateUser(u *User) error {
	return s.db.Save(u).Error
}

func (s *Store) DeleteUser(id uint, currentUserID uint) error {
	if id == currentUserID {
		return fmt.Errorf("cannot delete yourself")
	}
	user, err := s.GetUserByID(id)
	if err != nil {
		return err
	}
	if user.Role == "admin" {
		var adminCount int64
		s.db.Model(&User{}).Where("role = ?", "admin").Count(&adminCount)
		if adminCount <= 1 {
			return fmt.Errorf("cannot delete the last admin")
		}
	}
	_ = s.DeleteVisitorAccessLogs(id)
	return s.db.Delete(&User{}, id).Error
}

func (s *Store) RecordRequestIP(resourceType string, resourceID uint, visitorIP, directIP string) error {
	var rec RequestIP
	err := s.db.Where("resource_type = ? AND resource_id = ? AND ip = ?",
		resourceType, resourceID, visitorIP).First(&rec).Error
	now := time.Now()
	if err == gorm.ErrRecordNotFound {
		rec = RequestIP{
			ResourceType: resourceType,
			ResourceID:   resourceID,
			IP:           visitorIP,
			DirectIP:     directIP,
			HitCount:     1,
			LastSeenAt:   now,
		}
		return s.db.Create(&rec).Error
	}
	if err != nil {
		return err
	}
	return s.db.Model(&rec).Updates(map[string]interface{}{
		"hit_count":    rec.HitCount + 1,
		"last_seen_at": now,
		"direct_ip":    directIP,
	}).Error
}

func (s *Store) ListRequestIPs(resourceType string, resourceID uint, limit, offset int) ([]RequestIP, int64, error) {
	var items []RequestIP
	var total int64
	q := s.db.Model(&RequestIP{})
	if resourceType != "" {
		q = q.Where("resource_type = ?", resourceType)
	}
	if resourceID > 0 {
		q = q.Where("resource_id = ?", resourceID)
	}
	q.Count(&total)
	err := q.Order("last_seen_at desc").Limit(limit).Offset(offset).Find(&items).Error
	return items, total, err
}

func (s *Store) UpsertFlowStat(resourceType string, resourceID uint, inlet, export int64) error {
	var stat FlowStat
	err := s.db.Where("resource_type = ? AND resource_id = ?", resourceType, resourceID).First(&stat).Error
	if err == gorm.ErrRecordNotFound {
		stat = FlowStat{
			ResourceType: resourceType,
			ResourceID:   resourceID,
			InletBytes:   inlet,
			ExportBytes:  export,
		}
		return s.db.Create(&stat).Error
	}
	if err != nil {
		return err
	}
	return s.db.Model(&stat).Updates(map[string]interface{}{
		"inlet_bytes":  gorm.Expr("inlet_bytes + ?", inlet),
		"export_bytes": gorm.Expr("export_bytes + ?", export),
	}).Error
}

func (s *Store) AddClientFlow(id uint, inlet, export int64) error {
	return s.db.Model(&Client{}).Where("id = ?", id).Updates(map[string]interface{}{
		"inlet_flow":  gorm.Expr("inlet_flow + ?", inlet),
		"export_flow": gorm.Expr("export_flow + ?", export),
	}).Error
}

func (s *Store) AddTunnelFlow(id uint, inlet, export int64) error {
	return s.db.Model(&Tunnel{}).Where("id = ?", id).Updates(map[string]interface{}{
		"inlet_flow":  gorm.Expr("inlet_flow + ?", inlet),
		"export_flow": gorm.Expr("export_flow + ?", export),
	}).Error
}

func (s *Store) AddHostFlow(id uint, inlet, export int64) error {
	return s.db.Model(&Host{}).Where("id = ?", id).Updates(map[string]interface{}{
		"inlet_flow":  gorm.Expr("inlet_flow + ?", inlet),
		"export_flow": gorm.Expr("export_flow + ?", export),
	}).Error
}

func (s *Store) DashboardStats() (map[string]interface{}, error) {
	var clientTotal, clientOnline int64
	var tunnelTotal, tunnelRunning int64
	var hostTotal int64
	var inletFlow, exportFlow int64

	s.db.Model(&Client{}).Count(&clientTotal)
	s.db.Model(&Tunnel{}).Count(&tunnelTotal)
	s.db.Model(&Tunnel{}).Where("run_status = ?", true).Count(&tunnelRunning)
	s.db.Model(&Host{}).Count(&hostTotal)
	s.db.Model(&Client{}).Select("COALESCE(SUM(inlet_flow),0)").Scan(&inletFlow)
	s.db.Model(&Client{}).Select("COALESCE(SUM(export_flow),0)").Scan(&exportFlow)

	return map[string]interface{}{
		"client_total":   clientTotal,
		"client_online":  clientOnline,
		"tunnel_total":   tunnelTotal,
		"tunnel_running": tunnelRunning,
		"host_total":     hostTotal,
		"inlet_flow":     inletFlow,
		"export_flow":    exportFlow,
	}, nil
}

func (s *Store) TopClientsByFlow(limit int) ([]Client, error) {
	var clients []Client
	err := s.db.Order("(inlet_flow + export_flow) desc").Limit(limit).Find(&clients).Error
	return clients, err
}

func (s *Store) TopTunnelsByFlow(limit int) ([]Tunnel, error) {
	var tunnels []Tunnel
	err := s.db.Order("(inlet_flow + export_flow) desc").Limit(limit).Find(&tunnels).Error
	return tunnels, err
}

func (s *Store) TopHostsByFlow(limit int) ([]Host, error) {
	var hosts []Host
	err := s.db.Order("(inlet_flow + export_flow) desc").Limit(limit).Find(&hosts).Error
	return hosts, err
}

func (s *Store) TopRequestIPs(limit int) ([]RequestIP, error) {
	var items []RequestIP
	err := s.db.Order("hit_count desc").Limit(limit).Find(&items).Error
	return items, err
}

func ClientOwnerCheck(s *Store, clientID, userID uint, isAdmin bool) error {
	if isAdmin {
		return nil
	}
	c, err := s.GetClientByID(clientID)
	if err != nil {
		return err
	}
	if c.OwnerUserID != userID {
		return fmt.Errorf("permission denied")
	}
	return nil
}