Prism/internal/store/traffic_repo.go

package store

import (
	"context"
	"database/sql"
	"fmt"
)

func (s *Store) InsertProxyRequestLog(ctx context.Context, l *ProxyRequestLog) error {
	res, err := s.db.ExecContext(ctx, `
		INSERT INTO proxy_request_logs(host,network,rule_type,rule_payload,rule_line,outbound_node,outbound_chain,upload,download,success,started_at,ended_at)
		VALUES(?,?,?,?,?,?,?,?,?,?,?,?)`,
		l.Host, l.Network, l.RuleType, l.RulePayload, l.RuleLine, l.OutboundNode, l.OutboundChain,
		l.Upload, l.Download, intFromBool(l.Success), l.StartedAt.Format("2006-01-02 15:04:05"), l.EndedAt.Format("2006-01-02 15:04:05"))
	if err != nil {
		return err
	}
	id, _ := res.LastInsertId()
	l.ID = id
	return nil
}

func (s *Store) ListProxyRequestLogs(ctx context.Context, keyword string, limit, offset int) ([]ProxyRequestLog, int, error) {
	query := `SELECT id,host,network,rule_type,rule_payload,rule_line,outbound_node,outbound_chain,upload,download,success,started_at,ended_at,created_at FROM proxy_request_logs WHERE 1=1`
	countQuery := `SELECT COUNT(*) FROM proxy_request_logs WHERE 1=1`
	args := []any{}
	if keyword != "" {
		query += ` AND (host LIKE ? OR rule_line LIKE ? OR outbound_node LIKE ? OR outbound_chain LIKE ?)`
		countQuery += ` AND (host LIKE ? OR rule_line LIKE ? OR outbound_node LIKE ? OR outbound_chain LIKE ?)`
		kw := "%" + keyword + "%"
		args = append(args, kw, kw, kw, kw)
	}
	var total int
	if err := s.db.QueryRowContext(ctx, countQuery, args...).Scan(&total); err != nil {
		return nil, 0, err
	}
	query += ` ORDER BY id DESC LIMIT ? OFFSET ?`
	args = append(args, limit, offset)
	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()
	var out []ProxyRequestLog
	for rows.Next() {
		var l ProxyRequestLog
		var success int
		var started, ended, created string
		if err := rows.Scan(&l.ID, &l.Host, &l.Network, &l.RuleType, &l.RulePayload, &l.RuleLine, &l.OutboundNode, &l.OutboundChain,
			&l.Upload, &l.Download, &success, &started, &ended, &created); err != nil {
			return nil, 0, err
		}
		l.Success = boolFromInt(success)
		l.StartedAt = parseTime(started)
		l.EndedAt = parseTime(ended)
		l.CreatedAt = parseTime(created)
		out = append(out, l)
	}
	return out, total, rows.Err()
}

func (s *Store) PruneProxyRequestLogs(ctx context.Context, keep int) error {
	_, err := s.db.ExecContext(ctx, `DELETE FROM proxy_request_logs WHERE id NOT IN (SELECT id FROM proxy_request_logs ORDER BY id DESC LIMIT ?)`, keep)
	return err
}

type TrafficAgg struct {
	Key         string `json:"key"`
	Upload      int64  `json:"upload"`
	Download    int64  `json:"download"`
	Connections int    `json:"connections"`
}

type TrafficHostNodeAgg struct {
	Host        string `json:"host"`
	NodeKey     string `json:"node_key"`
	Upload      int64  `json:"upload"`
	Download    int64  `json:"download"`
	Connections int    `json:"connections"`
}

func (s *Store) AggregateTrafficByHost(ctx context.Context, hours, limit int) ([]TrafficAgg, error) {
	if hours <= 0 {
		hours = 24
	}
	if limit <= 0 {
		limit = 30
	}
	rows, err := s.db.QueryContext(ctx, `
		SELECT host, COALESCE(SUM(upload),0), COALESCE(SUM(download),0), COUNT(*)
		FROM proxy_request_logs
		WHERE host != '' AND created_at >= datetime('now', ?)
		GROUP BY host
		ORDER BY (SUM(upload)+SUM(download)) DESC
		LIMIT ?`, fmt.Sprintf("-%d hours", hours), limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanTrafficAgg(rows)
}

func (s *Store) CountAggregateTrafficByNode(ctx context.Context, hours int) (int, error) {
	if hours <= 0 {
		hours = 24
	}
	var n int
	err := s.db.QueryRowContext(ctx, `
		SELECT COUNT(*) FROM (
			SELECT 1 FROM proxy_request_logs
			WHERE outbound_node != '' AND created_at >= datetime('now', ?)
			GROUP BY outbound_node
		)`, fmt.Sprintf("-%d hours", hours)).Scan(&n)
	return n, err
}

func (s *Store) AggregateTrafficByNode(ctx context.Context, hours, limit, offset int) ([]TrafficAgg, error) {
	if hours <= 0 {
		hours = 24
	}
	if limit <= 0 {
		limit = 50
	}
	if offset < 0 {
		offset = 0
	}
	rows, err := s.db.QueryContext(ctx, `
		SELECT outbound_node, COALESCE(SUM(upload),0), COALESCE(SUM(download),0), COUNT(*)
		FROM proxy_request_logs
		WHERE outbound_node != '' AND created_at >= datetime('now', ?)
		GROUP BY outbound_node
		ORDER BY (SUM(upload)+SUM(download)) DESC
		LIMIT ? OFFSET ?`, fmt.Sprintf("-%d hours", hours), limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanTrafficAgg(rows)
}

func (s *Store) CountAggregateTrafficByHostNode(ctx context.Context, hours int) (int, error) {
	if hours <= 0 {
		hours = 24
	}
	var n int
	err := s.db.QueryRowContext(ctx, `
		SELECT COUNT(*) FROM (
			SELECT 1 FROM proxy_request_logs
			WHERE host != '' AND outbound_node != '' AND created_at >= datetime('now', ?)
			GROUP BY host, outbound_node
		)`, fmt.Sprintf("-%d hours", hours)).Scan(&n)
	return n, err
}

func (s *Store) AggregateTrafficByHostNode(ctx context.Context, hours, limit, offset int) ([]TrafficHostNodeAgg, error) {
	if hours <= 0 {
		hours = 24
	}
	if limit <= 0 {
		limit = 50
	}
	if offset < 0 {
		offset = 0
	}
	rows, err := s.db.QueryContext(ctx, `
		SELECT host, outbound_node, COALESCE(SUM(upload),0), COALESCE(SUM(download),0), COUNT(*)
		FROM proxy_request_logs
		WHERE host != '' AND outbound_node != '' AND created_at >= datetime('now', ?)
		GROUP BY host, outbound_node
		ORDER BY (SUM(upload)+SUM(download)) DESC
		LIMIT ? OFFSET ?`, fmt.Sprintf("-%d hours", hours), limit, offset)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []TrafficHostNodeAgg
	for rows.Next() {
		var item TrafficHostNodeAgg
		if err := rows.Scan(&item.Host, &item.NodeKey, &item.Upload, &item.Download, &item.Connections); err != nil {
			return nil, err
		}
		out = append(out, item)
	}
	return out, rows.Err()
}

func scanTrafficAgg(rows *sql.Rows) ([]TrafficAgg, error) {
	var out []TrafficAgg
	for rows.Next() {
		var item TrafficAgg
		if err := rows.Scan(&item.Key, &item.Upload, &item.Download, &item.Connections); err != nil {
			return nil, err
		}
		out = append(out, item)
	}
	return out, rows.Err()
}