Skip to content

Commit

Permalink
🐶 重构部分用户数据库操作代码
Browse files Browse the repository at this point in the history
  • Loading branch information
dushixiang committed Mar 17, 2021
1 parent e1cd732 commit 0150361
Show file tree
Hide file tree
Showing 50 changed files with 456 additions and 431 deletions.
245 changes: 118 additions & 127 deletions main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,18 +4,19 @@ import (
"bytes"
"fmt"
"io"
"next-terminal/server/repository"
"os"
"strconv"
"strings"
"time"

"next-terminal/pkg/api"
"next-terminal/pkg/config"
"next-terminal/pkg/constant"
"next-terminal/pkg/global"
"next-terminal/pkg/handle"
"next-terminal/pkg/model"
"next-terminal/pkg/utils"
"next-terminal/server/api"
"next-terminal/server/config"
"next-terminal/server/constant"
"next-terminal/server/global"
"next-terminal/server/handle"
"next-terminal/server/model"
"next-terminal/server/utils"

nested "github.com/antonfisher/nested-logrus-formatter"
"github.com/labstack/gommon/log"
Expand All @@ -30,6 +31,11 @@ import (

const Version = "v0.3.3"

var (
db *gorm.DB
userRepository repository.UserRepository
)

func main() {
err := Run()
if err != nil {
Expand Down Expand Up @@ -64,66 +70,67 @@ func Run() error {

logrus.SetOutput(io.MultiWriter(writer1, writer2, writer3))

global.Config, err = config.SetupConfig()
if err != nil {
return err
global.Config = config.SetupConfig()
db = SetupDB()

// 初始化 repository
global.DB = db
userRepository = repository.UserRepository{DB: db}

if global.Config.ResetPassword != "" {
return ResetPassword()
}

var logMode logger.Interface
if global.Config.Debug {
logMode = logger.Default.LogMode(logger.Info)
} else {
logMode = logger.Default.LogMode(logger.Silent)
if err := global.DB.AutoMigrate(&model.User{}, &model.Asset{}, &model.AssetAttribute{}, &model.Session{}, &model.Command{},
&model.Credential{}, &model.Property{}, &model.ResourceSharer{}, &model.UserGroup{}, &model.UserGroupMember{},
&model.LoginLog{}, &model.Num{}, &model.Job{}, &model.JobLog{}, &model.AccessSecurity{}); err != nil {
return err
}

fmt.Printf("当前数据库模式为：%v\n", global.Config.DB)
if global.Config.DB == "mysql" {
dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Local",
global.Config.Mysql.Username,
global.Config.Mysql.Password,
global.Config.Mysql.Hostname,
global.Config.Mysql.Port,
global.Config.Mysql.Database,
)
global.DB, err = gorm.Open(mysql.Open(dsn), &gorm.Config{
Logger: logMode,
})
} else {
global.DB, err = gorm.Open(sqlite.Open(global.Config.Sqlite.File), &gorm.Config{
Logger: logMode,
})
if err := InitDBData(); err != nil {
return err
}

if err != nil {
logrus.Errorf("连接数据库异常：%v", err.Error())
if err := api.ReloadAccessSecurity(); err != nil {
return err
}

if global.Config.ResetPassword != "" {
user, err := model.FindUserByUsername(global.Config.ResetPassword)
if err != nil {
return err
}
password := "next-terminal"
passwd, err := utils.Encoder.Encode([]byte(password))
if err != nil {
return err
}
u := &model.User{
Password: string(passwd),
// 配置缓存器
global.Cache = cache.New(5*time.Minute, 10*time.Minute)
global.Cache.OnEvicted(func(key string, value interface{}) {
if strings.HasPrefix(key, api.Token) {
token := api.GetTokenFormCacheKey(key)
logrus.Debugf("用户Token「%v」过期", token)
err := model.Logout(token)
if err != nil {
logrus.Errorf("退出登录失败 %v", err)
}
}
model.UpdateUserById(u, user.ID)
logrus.Debugf("用户「%v」密码初始化为: %v", user.Username, password)
return nil
}
})
global.Store = global.NewStore()
global.Cron = cron.New(cron.WithSeconds()) //精确到秒
global.Cron.Start()

if err := global.DB.AutoMigrate(&model.User{}); err != nil {
e := api.SetupRoutes(userRepository)
if err := handle.InitProperties(); err != nil {
return err
}
// 启动定时任务
go handle.RunTicker()
go handle.RunDataFix()

users := model.FindAllUser()
if len(users) == 0 {
if global.Config.Server.Cert != "" && global.Config.Server.Key != "" {
return e.StartTLS(global.Config.Server.Addr, global.Config.Server.Cert, global.Config.Server.Key)
} else {
return e.Start(global.Config.Server.Addr)
}

}

func InitDBData() (err error) {
users := userRepository.FindAll()

if len(users) == 0 {
initPassword := "admin"
var pass []byte
if pass, err = utils.Encoder.Encode([]byte(initPassword)); err != nil {
Expand All @@ -138,7 +145,7 @@ func Run() error {
Type: constant.TypeAdmin,
Created: utils.NowJsonTime(),
}
if err := model.CreateNewUser(&user); err != nil {
if err := userRepository.Create(&user); err != nil {
return err
}
logrus.Infof("初始用户创建成功，账号：「%v」密码：「%v」", user.Username, initPassword)
Expand All @@ -148,59 +155,16 @@ func Run() error {
if users[i].Type == "" {
user := model.User{
Type: constant.TypeAdmin,
ID: users[i].ID,
}
if err := userRepository.Update(&user); err != nil {
return err
}
model.UpdateUserById(&user, users[i].ID)
logrus.Infof("自动修正用户「%v」ID「%v」类型为管理员", users[i].Nickname, users[i].ID)
}
}
}

if err := global.DB.AutoMigrate(&model.Asset{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.AssetAttribute{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.Session{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.Command{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.Credential{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.Property{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.ResourceSharer{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.UserGroup{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.UserGroupMember{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.LoginLog{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.Num{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.Job{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.JobLog{}); err != nil {
return err
}
if err := global.DB.AutoMigrate(&model.AccessSecurity{}); err != nil {
return err
}
if err := api.ReloadAccessSecurity(); err != nil {
return err
}

if len(model.FindAllTemp()) == 0 {
for i := 0; i <= 30; i++ {
if err := model.CreateNewTemp(&model.Num{I: strconv.Itoa(i)}); err != nil {
Expand All @@ -209,22 +173,6 @@ func Run() error {
}
}

// 配置缓存器
global.Cache = cache.New(5*time.Minute, 10*time.Minute)
global.Cache.OnEvicted(func(key string, value interface{}) {
if strings.HasPrefix(key, api.Token) {
token := api.GetTokenFormCacheKey(key)
logrus.Debugf("用户Token「%v」过期", token)
err := model.Logout(token)
if err != nil {
logrus.Errorf("退出登录失败 %v", err)
}
}
})
global.Store = global.NewStore()
global.Cron = cron.New(cron.WithSeconds()) //精确到秒
global.Cron.Start()

jobs, err := model.FindJobByFunc(constant.FuncCheckAssetStatusJob)
if err != nil {
return err
Expand Down Expand Up @@ -264,7 +212,7 @@ func Run() error {
for i := range loginLogs {
loginLog := loginLogs[i]
token := loginLog.ID
user, err := model.FindUserById(loginLog.UserId)
user, err := userRepository.FindById(loginLog.UserId)
if err != nil {
logrus.Debugf("用户「%v」获取失败，忽略", loginLog.UserId)
continue
Expand All @@ -288,7 +236,7 @@ func Run() error {
}

// 修正用户登录状态
onlineUsers, err := model.FindOnlineUsers()
onlineUsers, err := userRepository.FindOnlineUsers()
if err != nil {
return err
}
Expand All @@ -298,24 +246,67 @@ func Run() error {
return err
}
if len(logs) == 0 {
if err := model.UpdateUserOnline(false, onlineUsers[i].ID); err != nil {
if err := userRepository.UpdateOnline(onlineUsers[i].ID, false); err != nil {
return err
}
}
}

e := api.SetupRoutes()
if err := handle.InitProperties(); err != nil {
return nil
}

func ResetPassword() error {
user, err := userRepository.FindByUsername(global.Config.ResetPassword)
if err != nil {
return err
}
// 启动定时任务
go handle.RunTicker()
go handle.RunDataFix()
password := "next-terminal"
passwd, err := utils.Encoder.Encode([]byte(password))
if err != nil {
return err
}
u := &model.User{
Password: string(passwd),
ID: user.ID,
}
if err := userRepository.Update(u); err != nil {
return err
}
logrus.Debugf("用户「%v」密码初始化为: %v", user.Username, password)
return nil
}

if global.Config.Server.Cert != "" && global.Config.Server.Key != "" {
return e.StartTLS(global.Config.Server.Addr, global.Config.Server.Cert, global.Config.Server.Key)
func SetupDB() *gorm.DB {

var logMode logger.Interface
if global.Config.Debug {
logMode = logger.Default.LogMode(logger.Info)
} else {
return e.Start(global.Config.Server.Addr)
logMode = logger.Default.LogMode(logger.Silent)
}

fmt.Printf("当前数据库模式为：%v\n", global.Config.DB)
var err error
var db *gorm.DB
if global.Config.DB == "mysql" {
dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4&parseTime=True&loc=Local",
global.Config.Mysql.Username,
global.Config.Mysql.Password,
global.Config.Mysql.Hostname,
global.Config.Mysql.Port,
global.Config.Mysql.Database,
)
db, err = gorm.Open(mysql.Open(dsn), &gorm.Config{
Logger: logMode,
})
} else {
db, err = gorm.Open(sqlite.Open(global.Config.Sqlite.File), &gorm.Config{
Logger: logMode,
})
}

if err != nil {
logrus.WithError(err).Panic("连接数据库异常")
}
return db
}
Loading

0 comments on commit 0150361

Please sign in to comment.