From 95b0f190ff6fe88dd12d51eedc2e4d71d8cf3715 Mon Sep 17 00:00:00 2001
From: GUO YANKE <hi@yankeguo.com>
Date: Tue, 23 Jan 2024 12:10:03 +0800
Subject: [PATCH] feat: add initializeAdminUsers

---
 .gitignore  |  9 ++++-----
 database.go | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 go.mod      |  2 +-
 main.go     |  1 +
 4 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/.gitignore b/.gitignore
index 186a2c1..a6553c0 100644
--- a/.gitignore
+++ b/.gitignore
@@ -20,13 +20,12 @@
 # Go workspace file
 go.work
 
-# Binary and Config
+# Binary
 /bunker
-/bunker.yaml
 
-# Database
+# Data
+/bunker.yaml
 /database.sqlite3
-
-# SSH Keys
+/admin.yaml
 /ssh_host_*_key
 /ssh_client_*_key
diff --git a/database.go b/database.go
index a3ece5d..56dd14c 100644
--- a/database.go
+++ b/database.go
@@ -1,13 +1,68 @@
 package main
 
 import (
+	"bytes"
+	"errors"
+	"os"
 	"path/filepath"
+	"time"
 
 	"github.com/glebarez/sqlite"
 	"github.com/yankeguo/bunker/model"
+	"github.com/yankeguo/bunker/model/dao"
+	"gopkg.in/yaml.v3"
 	"gorm.io/gorm"
 )
 
+func initializeAdminUsers(dataDir DataDir, _db *gorm.DB) (err error) {
+	buf, _ := os.ReadFile(filepath.Join(dataDir.String(), "admin.yaml"))
+	buf = bytes.TrimSpace(buf)
+	if len(buf) == 0 {
+		return
+	}
+
+	var m map[string]string
+	if err = yaml.Unmarshal(buf, &m); err != nil {
+		return
+	}
+
+	db := dao.Use(_db)
+
+	for username, password := range m {
+		var user *model.User
+		if user, err = db.User.Where(db.User.ID.Eq(username)).First(); err != nil {
+			if errors.Is(err, gorm.ErrRecordNotFound) {
+				user = &model.User{
+					ID:        username,
+					CreatedAt: time.Now(),
+					VisitedAt: time.Now(),
+					IsAdmin:   true,
+				}
+				user.SetPassword(password)
+
+				if err = db.User.Create(user); err != nil {
+					return
+				}
+			} else {
+				return
+			}
+		} else {
+			user.SetPassword(password)
+
+			if _, err = db.User.Where(
+				db.User.ID.Eq(username),
+			).UpdateSimple(
+				db.User.PasswordDigest.Value(user.PasswordDigest),
+				db.User.IsAdmin.Value(true),
+			); err != nil {
+				return
+			}
+		}
+	}
+
+	return
+}
+
 func createDatabase(dataDir DataDir) (db *gorm.DB, err error) {
 	if db, err = gorm.Open(
 		sqlite.Open(filepath.Join(dataDir.String(), "database.sqlite3")),
@@ -15,7 +70,7 @@ func createDatabase(dataDir DataDir) (db *gorm.DB, err error) {
 	); err != nil {
 		return
 	}
-	if err = db.Debug().AutoMigrate(model.All...); err != nil {
+	if err = db.AutoMigrate(model.All...); err != nil {
 		return
 	}
 	return
diff --git a/go.mod b/go.mod
index dc441b9..9c4a7b3 100644
--- a/go.mod
+++ b/go.mod
@@ -9,6 +9,7 @@ require (
 	github.com/yankeguo/ufx v0.2.2
 	go.uber.org/fx v1.20.1
 	golang.org/x/crypto v0.18.0
+	gopkg.in/yaml.v3 v3.0.1
 	gorm.io/gen v0.3.25
 	gorm.io/gorm v1.25.5
 	gorm.io/plugin/dbresolver v1.5.0
@@ -56,7 +57,6 @@ require (
 	golang.org/x/text v0.14.0 // indirect
 	golang.org/x/tools v0.17.0 // indirect
 	google.golang.org/protobuf v1.32.0 // indirect
-	gopkg.in/yaml.v3 v3.0.1 // indirect
 	gorm.io/datatypes v1.2.0 // indirect
 	gorm.io/driver/mysql v1.5.2 // indirect
 	gorm.io/hints v1.1.2 // indirect
diff --git a/main.go b/main.go
index cd9f068..7914606 100644
--- a/main.go
+++ b/main.go
@@ -30,6 +30,7 @@ func main() {
 		ufx.Module,
 		fx.Invoke(installStatic, installAPIAuthorizedKeys),
 		fx.Invoke(func(s *SSHServer) {}),
+		fx.Invoke(initializeAdminUsers),
 	)
 	if app.Err() != nil {
 		log.Println(app.Err().Error())