modify main.go 1

This commit is contained in:
2025-03-26 02:36:29 +08:00
parent 29a5bfe9d9
commit 0c77e0a28d

View File

@ -17,6 +17,7 @@ import (
"gopkg.in/yaml.v2"
)
// 配置结构体
type Config struct {
Server struct {
Port string `yaml:"port"`
@ -31,25 +32,29 @@ type Config struct {
} `yaml:"admin"`
}
// 网站结构体
type Site struct {
ID int `json:"id"`
Name string `json:"name"`
URL string `json:"url"`
Icon string `json:"icon,omitempty"`
Category string `json:"category"`
Order int `json:"order"` // 新增排序字段
Order int `json:"order"`
}
// 登录凭据
type Credentials struct {
Username string `json:"username"`
Password string `json:"password"`
}
// JWT声明
type Claims struct {
Username string `json:"username"`
jwt.StandardClaims
}
// 全局变量
var (
db *sql.DB
config *Config
@ -72,6 +77,7 @@ func main() {
// 设置路由
http.HandleFunc("/api/login", loginHandler)
http.HandleFunc("/api/validate", authMiddleware(validateHandler))
http.HandleFunc("/api/sites", sitesHandler)
http.HandleFunc("/api/sites/order", authMiddleware(updateSitesOrderHandler))
http.HandleFunc("/api/sites/add", authMiddleware(addSiteHandler))
@ -84,10 +90,11 @@ func main() {
http.Handle("/", fs)
// 启动服务器
log.Printf("Starting server on :%s", config.Server.Port)
log.Printf("Server starting on :%s", config.Server.Port)
log.Fatal(http.ListenAndServe(":"+config.Server.Port, nil))
}
// 加载配置文件
func loadConfig(path string) (*Config, error) {
cfg := &Config{}
file, err := os.Open(path)
@ -104,6 +111,7 @@ func loadConfig(path string) (*Config, error) {
return cfg, nil
}
// 初始化数据库
func initDB(dbPath string) error {
var err error
os.MkdirAll(filepath.Dir(dbPath), 0755)
@ -113,7 +121,7 @@ func initDB(dbPath string) error {
return err
}
// 创建表
// 创建网站
_, err = db.Exec(`CREATE TABLE IF NOT EXISTS sites (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL,
@ -137,13 +145,17 @@ func initDB(dbPath string) error {
}
// 初始化管理员账户
hashedPassword, _ := bcrypt.GenerateFromPassword([]byte(config.Admin.Password), bcrypt.DefaultCost)
db.Exec("INSERT OR IGNORE INTO admin_users (username, password) VALUES (?, ?)",
config.Admin.Username, string(hashedPassword))
return nil
hashedPassword, err := bcrypt.GenerateFromPassword([]byte(config.Admin.Password), bcrypt.DefaultCost)
if err != nil {
return err
}
_, err = db.Exec("INSERT OR IGNORE INTO admin_users (username, password) VALUES (?, ?)",
config.Admin.Username, string(hashedPassword))
return err
}
// 登录处理器
func loginHandler(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
@ -164,6 +176,7 @@ func loginHandler(w http.ResponseWriter, r *http.Request) {
return
}
// 验证密码
if err := bcrypt.CompareHashAndPassword([]byte(storedPassword), []byte(creds.Password)); err != nil {
http.Error(w, "Invalid credentials", http.StatusUnauthorized)
return
@ -192,6 +205,15 @@ func loginHandler(w http.ResponseWriter, r *http.Request) {
})
}
// 验证处理器
func validateHandler(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]bool{
"valid": true,
})
}
// 认证中间件
func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
authHeader := r.Header.Get("Authorization")
@ -207,7 +229,16 @@ func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
return []byte(config.Server.SecretKey), nil
})
if err != nil || !token.Valid {
if err != nil {
if err == jwt.ErrSignatureInvalid {
http.Error(w, "Invalid token signature", http.StatusUnauthorized)
return
}
http.Error(w, "Invalid token", http.StatusUnauthorized)
return
}
if !token.Valid {
http.Error(w, "Invalid token", http.StatusUnauthorized)
return
}
@ -216,6 +247,7 @@ func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
}
}
// 获取所有网站
func sitesHandler(w http.ResponseWriter, r *http.Request) {
if r.Method != "GET" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
@ -243,6 +275,7 @@ func sitesHandler(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(sites)
}
// 更新网站排序
func updateSitesOrderHandler(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
@ -278,6 +311,7 @@ func updateSitesOrderHandler(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}
// 添加网站
func addSiteHandler(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
@ -297,8 +331,8 @@ func addSiteHandler(w http.ResponseWriter, r *http.Request) {
}
// 插入数据库
res, err := db.Exec("INSERT INTO sites (name, url, icon, category) VALUES (?, ?, ?, ?)",
s.Name, s.URL, s.Icon, s.Category)
res, err := db.Exec("INSERT INTO sites (name, url, icon, category, sort_order) VALUES (?, ?, ?, ?, ?)",
s.Name, s.URL, s.Icon, s.Category, s.Order)
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
@ -312,6 +346,7 @@ func addSiteHandler(w http.ResponseWriter, r *http.Request) {
json.NewEncoder(w).Encode(s)
}
// 删除网站
func deleteSiteHandler(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
@ -333,5 +368,81 @@ func deleteSiteHandler(w http.ResponseWriter, r *http.Request) {
}
w.WriteHeader(http.StatusOK)
fmt.Fprint(w, `{"status": "success"}`)
json.NewEncoder(w).Encode(map[string]string{"status": "success"})
}
// 导出数据
func exportHandler(w http.ResponseWriter, r *http.Request) {
if r.Method != "GET" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
rows, err := db.Query("SELECT id, name, url, icon, category, sort_order FROM sites ORDER BY category, sort_order, name")
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
defer rows.Close()
var sites []Site
for rows.Next() {
var s Site
if err := rows.Scan(&s.ID, &s.Name, &s.URL, &s.Icon, &s.Category, &s.Order); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
sites = append(sites, s)
}
w.Header().Set("Content-Type", "application/json")
w.Header().Set("Content-Disposition", "attachment; filename=navigation-sites.json")
json.NewEncoder(w).Encode(sites)
}
// 导入数据
func importHandler(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" {
http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
return
}
var sites []Site
if err := json.NewDecoder(r.Body).Decode(&sites); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
tx, err := db.Begin()
if err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
// 清空现有数据
_, err = tx.Exec("DELETE FROM sites")
if err != nil {
tx.Rollback()
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
// 插入新数据
for _, site := range sites {
_, err = tx.Exec("INSERT INTO sites (name, url, icon, category, sort_order) VALUES (?, ?, ?, ?, ?)",
site.Name, site.URL, site.Icon, site.Category, site.Order)
if err != nil {
tx.Rollback()
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
}
if err := tx.Commit(); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
return
}
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(map[string]string{"status": "success"})
}