Skip to content
This repository has been archived by the owner on Feb 5, 2023. It is now read-only.

Commit

Permalink
Add API admin features
Browse files Browse the repository at this point in the history
Antonio Cheong committed Dec 27, 2022
1 parent 6f564c2 commit 53a1de8
Showing 8 changed files with 268 additions and 5 deletions.
Binary file added Data/auth.db
Binary file not shown.
Empty file added Data/placeholder
Empty file.
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
@@ -6,6 +6,7 @@ require (
github.com/gin-gonic/gin v1.8.1
github.com/google/uuid v1.3.0
github.com/gorilla/websocket v1.5.0
github.com/mattn/go-sqlite3 v1.14.16
)

require (
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
@@ -38,6 +38,8 @@ github.com/leodido/go-urn v1.2.1 h1:BqpAaACuzVSgi/VLzGZIobT2z4v53pjosyNd9Yv6n/w=
github.com/leodido/go-urn v1.2.1/go.mod h1:zt4jvISO2HfUBqxjfIshjdMTYS56ZS/qv49ictyFfxY=
github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y=
github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
github.com/mattn/go-sqlite3 v1.14.16 h1:yOQRA0RpS5PFz/oikGwBEqvAWhWg5ufRz4ETLjwpU1Y=
github.com/mattn/go-sqlite3 v1.14.16/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421 h1:ZqeYNhU3OHLH3mGKHDcjJRFFRrJa6eAM5H+CtDdOsPc=
github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M=
116 changes: 116 additions & 0 deletions handlers/admin.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
package handlers

import (
"os"

"github.com/ChatGPT-Hackers/ChatGPT-API-server/utils"
_ "github.com/ChatGPT-Hackers/ChatGPT-API-server/utils"
"github.com/gin-gonic/gin"
)

type Request struct {
AdminKey string `json:"admin_key"`
UserID string `json:"user_id"`
}

func Admin_userAdd(c *gin.Context) {
// Get admin key from request body
var request Request
if err := c.ShouldBindJSON(&request); err != nil {
c.JSON(400, gin.H{
"error": "Invalid request body",
})
return
}

// Check if admin key is valid
if !utils.VerifyAdminKey(request.AdminKey) {
c.JSON(401, gin.H{
"error": "Invalid admin key",
})
return
}

// Generate user_id and token
user_id := utils.GenerateId()
token := utils.GenerateId()

// Insert user_id and token into database
err := utils.DatabaseInsert(user_id, token)
if err != nil {
c.JSON(500, gin.H{
"error": "Failed to insert user_id and token into database",
})
return
}

// Return user_id and token
c.JSON(200, gin.H{
"user_id": user_id,
"token": token,
})
}

// POST request to delete a user
func Admin_userDel(c *gin.Context) {
// Get admin key from request body
var request Request
if err := c.ShouldBindJSON(&request); err != nil {
c.JSON(400, gin.H{
"error": "Invalid request body",
})
return
}

// Check if admin key is valid
if !utils.VerifyAdminKey(request.AdminKey) {
c.JSON(401, gin.H{
"error": "Invalid admin key",
})
return
}

// Delete user from database
err := utils.DatabaseDelete(request.UserID)
if err != nil {
c.JSON(500, gin.H{
"error": "Failed to delete user from database",
})
return
}

// Return success
c.JSON(200, gin.H{
"message": "User deleted",
})
}

func Admin_usersGet(c *gin.Context) {
// Get admin key from GET parameter
AdminKey := c.Query("admin_key")

// Check if admin key is valid
if !utils.VerifyAdminKey(AdminKey) {
c.JSON(401, gin.H{
"error": "Invalid admin key",
"key": AdminKey,
"correct": os.Args[2],
})
return
}

// Get users from database
users, err := utils.DatabaseSelectAll()
if err != nil {
c.JSON(500, gin.H{
"message": "Failed to get users from database",
"error": err.Error(),
})
return
}

// Return users
c.JSON(200, gin.H{
"users": users,
})
}
12 changes: 9 additions & 3 deletions handlers/api.go
Original file line number Diff line number Diff line change
@@ -2,7 +2,6 @@ package handlers

import (
"encoding/json"
"os"
"time"

"github.com/ChatGPT-Hackers/ChatGPT-API-server/types"
@@ -28,7 +27,15 @@ func API_ask(c *gin.Context) {
})
return
}
if c.Request.Header["Authorization"][0] != os.Args[2] {
// Check if API key is valid
verified, err := utils.VerifyToken(c.Request.Header["Authorization"][0])
if err != nil {
c.JSON(500, gin.H{
"error": "Failed to verify API key",
})
return
}
if !verified {
c.JSON(401, gin.H{
"error": "Invalid API key",
})
@@ -185,7 +192,6 @@ func API_getConnections(c *gin.Context) {
c.JSON(200, gin.H{
"connections": connections,
})
return
}

func ping(connection_id string) bool {
15 changes: 13 additions & 2 deletions main.go
Original file line number Diff line number Diff line change
@@ -9,17 +9,26 @@ import (
"os"

"github.com/ChatGPT-Hackers/ChatGPT-API-server/handlers"
"github.com/ChatGPT-Hackers/ChatGPT-API-server/utils"

"github.com/gin-gonic/gin"
)

func main() {
// get arg server port and secret key
// get arg server port and admin key
if len(os.Args) < 3 {
println("Usage: ./ChatGPT-API-server <port> <secret key>")
println("Usage: ./ChatGPT-API-server <port> <admin key>")
return
}
println(os.Args[1], os.Args[2])

// Make database
err := utils.DatabaseCreate()
if err != nil {
println("Failed to create database")
return
}

router := gin.Default()

//// # Headers
@@ -32,6 +41,8 @@ func main() {
router.GET("/client/register", handlers.Client_register)
router.POST("/api/ask", handlers.API_ask)
router.GET("/api/connections", handlers.API_getConnections)
router.POST("/admin/users/add", handlers.Admin_userAdd)
router.GET("/admin/users", handlers.Admin_usersGet)

// Add a health endpoint
router.GET("/health", func(c *gin.Context) {
127 changes: 127 additions & 0 deletions utils/auth.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
package utils

import (
"database/sql"
"fmt"
"os"

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

func DatabaseCreate() error {
// Open a connection to the SQLite database
db, err := sql.Open("sqlite3", "./Data/auth.db")
if err != nil {
return fmt.Errorf("error opening database: %v", err)
}
defer db.Close()

// Create the table if it doesn't already exist
_, err = db.Exec(`CREATE TABLE IF NOT EXISTS tokens (user_id TEXT PRIMARY KEY, token TEXT UNIQUE)`)
if err != nil {
return fmt.Errorf("error creating table: %v", err)
}

return nil
}

func DatabaseInsert(user_id string, token string) error {
// Open a connection to the SQLite database
db, err := sql.Open("sqlite3", "./Data/auth.db")
if err != nil {
return fmt.Errorf("error opening database: %v", err)
}
defer db.Close()

// Insert the token into the database
_, err = db.Exec(`INSERT INTO tokens (user_id, token) VALUES (?, ?)`, user_id, token)
if err != nil {
return fmt.Errorf("error inserting token: %v", err)
}

return nil
}

func DatabaseDelete(user_id string) error {
// Open a connection to the SQLite database
db, err := sql.Open("sqlite3", "./Data/auth.db")
if err != nil {
return fmt.Errorf("error opening database: %v", err)
}
defer db.Close()

// Delete the token from the database
_, err = db.Exec(`DELETE FROM tokens WHERE user_id = ?`, user_id)
if err != nil {
return fmt.Errorf("error deleting token: %v", err)
}

return nil
}

type User struct {
UserID string `json:"user_id"`
Token string `json:"token"`
}

func DatabaseSelectAll() ([]User, error) {
// Open a connection to the SQLite database
db, err := sql.Open("sqlite3", "./Data/auth.db")
if err != nil {
return nil, fmt.Errorf("error opening database: %v", err)
}
defer db.Close()

// Select the token from the database
rows, err := db.Query(`SELECT * FROM tokens`)
if err != nil {
return nil, fmt.Errorf("error selecting token: %v", err)
}
defer rows.Close()

// user map ({"users": ["user_id": "...", "token": "..."], ...})
var users []User
for rows.Next() {
var user User
err = rows.Scan(&user.UserID, &user.Token)
if err != nil {
return nil, fmt.Errorf("error scanning rows: %v", err)
}
users = append(users, user)
}

return users, nil

}

// Verify admin key
func VerifyAdminKey(key string) bool {
return key == os.Args[2]
}

func VerifyToken(token string) (bool, error) {
// Check if token is admin key
if VerifyAdminKey(token) {
return true, nil
}
// Open a connection to the SQLite database
db, err := sql.Open("sqlite3", "./Data/auth.db")
if err != nil {
return false, fmt.Errorf("error opening database: %v", err)
}
defer db.Close()

// Select the token from the database
rows, err := db.Query(`SELECT * FROM tokens WHERE token = ?`, token)
if err != nil {
return false, fmt.Errorf("error selecting token: %v", err)
}
defer rows.Close()

// Check if the token exists
if rows.Next() {
return true, nil
} else {
return false, nil
}
}

0 comments on commit 53a1de8

Please sign in to comment.