1 files changed, 102 insertions, 6 deletions

diff --git a/db.go b/db.go
index 86cedb3..8f7a5c7 100644
--- a/db.go
+++ b/db.go
@@ -8,22 +8,68 @@ import (
 	"go.rikki.moe/v2stat/command"
 )
 
+type TrafficDirection int
+
 const (
-	DirectionDownlink = iota
+	DirectionDownlink TrafficDirection = iota
 	DirectionUplink
 )
 
+type ConnectionType int
+
 const (
-	ConnTypeUser = iota
+	ConnTypeUser ConnectionType = iota
 	ConnTypeInbound
 	ConnTypeOutbound
 )
 
 type ConnInfo struct {
-	Type     int    `json:"type"`
-	Name     string `json:"name"`
+	Type ConnectionType `json:"type"`
+	Name string         `json:"name"`
+}
+
+type TrafficStat struct {
+	Time     string `json:"time"`
+	Downlink int64  `json:"downlink"`
+	Uplink   int64  `json:"uplink"`
+}
+
+func (ci *ConnInfo) String() string {
+	switch ci.Type {
+	case ConnTypeUser:
+		return "user:" + ci.Name
+	case ConnTypeInbound:
+		return "inbound:" + ci.Name
+	case ConnTypeOutbound:
+		return "outbound:" + ci.Name
+	default:
+		return "unknown:" + ci.Name
+	}
+}
+
+func ParseConnInfo(s string) (ConnInfo, bool) {
+	parts := strings.Split(s, ":")
+	if len(parts) != 2 {
+		return ConnInfo{}, false
+	}
+	var connType ConnectionType
+	switch parts[0] {
+	case "user":
+		connType = ConnTypeUser
+	case "inbound":
+		connType = ConnTypeInbound
+	case "outbound":
+		connType = ConnTypeOutbound
+	default:
+		return ConnInfo{}, false
+	}
+	return ConnInfo{
+		Type: connType,
+		Name: parts[1],
+	}, true
 }
 
+
 func (v *V2Stat) InitDB() error {
 	stmts := []string{
 		`CREATE TABLE IF NOT EXISTS conn (
@@ -129,7 +175,7 @@ func (v *V2Stat) RecordNow(ctx context.Context) error {
 	return nil
 }
 
-func parseStatKey(key string) (connType int, connName string, direction int, ok bool) {
+func parseStatKey(key string) (connType ConnectionType, connName string, direction TrafficDirection, ok bool) {
 	parts := strings.Split(key, ">>>")
 	if len(parts) != 4 || parts[2] != "traffic" {
 		return 0, "", 0, false
@@ -158,4 +204,54 @@ func parseStatKey(key string) (connType int, connName string, direction int, ok
 	}
 
 	return connType, connName, direction, true
-}
\ No newline at end of file
+}
+
+func (v *V2Stat) QueryConn() ([]ConnInfo, error) {
+	rows, err := v.db.Query("SELECT type, name FROM conn")
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var conns []ConnInfo
+	for rows.Next() {
+		var conn ConnInfo
+		if err := rows.Scan(&conn.Type, &conn.Name); err != nil {
+			return nil, err
+		}
+		conns = append(conns, conn)
+	}
+	return conns, nil
+}
+
+func (v *V2Stat) QueryStatsHourly(conn *ConnInfo) ([]TrafficStat, error) {
+	rows, err := v.db.Query(`
+	SELECT
+	  strftime('%Y-%m-%d %H:00:00', datetime(s.timestamp, 'unixepoch', '+8 hours')) AS time,
+	  SUM(CASE WHEN s.direction = 0 THEN s.traffic ELSE 0 END) AS downlink,
+	  SUM(CASE WHEN s.direction = 1 THEN s.traffic ELSE 0 END) AS uplink
+	FROM stats s
+	JOIN conn c ON s.conn_id = c.id
+	WHERE c.type = ? AND c.name = ?
+	GROUP BY time
+	ORDER BY time;
+`, conn.Type, conn.Name)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	var stats []TrafficStat
+	for rows.Next() {
+		var s TrafficStat
+		if err := rows.Scan(&s.Time, &s.Downlink, &s.Uplink); err != nil {
+			panic(err)
+		}
+		stats = append(stats, s)
+	}
+
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return stats, nil		
+}