diff --git a/config.tpl.yaml b/config.tpl.yaml index d616e14..ef4193d 100644 --- a/config.tpl.yaml +++ b/config.tpl.yaml @@ -8,6 +8,9 @@ http: ip_header: dns: enable: true +rmi: + enable: true + addr: :1099 mysql: enable: true addr: :3306 diff --git a/frontend/src/App.vue b/frontend/src/App.vue index eca246b..abbe6ca 100644 --- a/frontend/src/App.vue +++ b/frontend/src/App.vue @@ -17,6 +17,9 @@ DNS Logs + + RMI Logs + MySQL Logs @@ -29,6 +32,9 @@ DNS Rules + + RMI Rules + MySQL Rules diff --git a/frontend/src/api/record.js b/frontend/src/api/record.js index 9bccae0..00b53bd 100644 --- a/frontend/src/api/record.js +++ b/frontend/src/api/record.js @@ -31,4 +31,15 @@ export function getMysqlRecord(params) { return status >= 200 && status < 300 // 默认的 } }) +} + +export function getRmiRecord(params) { + return request({ + url: '/record/rmi', + params: params, + method: 'get', + validateStatus: function (status) { + return status >= 200 && status < 300 // 默认的 + } + }) } \ No newline at end of file diff --git a/frontend/src/api/rule.js b/frontend/src/api/rule.js index f83437f..d6933ff 100644 --- a/frontend/src/api/rule.js +++ b/frontend/src/api/rule.js @@ -98,4 +98,37 @@ export function deleteMysqlRule(data) { return status >= 200 && status < 300 // 默认的 } }) +} + +export function getRmiRule(params) { + return request({ + url: '/rule/rmi', + params: params, + method: 'get', + validateStatus: function (status) { + return status >= 200 && status < 300 // 默认的 + } + }) +} + +export function upsertRmiRule(data) { + return request({ + url: '/rule/rmi', + data: data, + method: 'post', + validateStatus: function (status) { + return status >= 200 && status < 300 // 默认的 + } + }) +} + +export function deleteRmiRule(data) { + return request({ + url: '/rule/rmi', + data: data, + method: 'delete', + validateStatus: function (status) { + return status >= 200 && status < 300 // 默认的 + } + }) } \ No newline at end of file diff --git a/frontend/src/router/index.js b/frontend/src/router/index.js index 590026e..4474505 100644 --- a/frontend/src/router/index.js +++ b/frontend/src/router/index.js @@ -25,6 +25,11 @@ const routes = [ name: 'MysqlLogs', component: () => import(/* webpackChunkName: "about" */ '../views/logs/Mysql') }, + { + path: '/logs/rmi', + name: 'RmiLogs', + component: () => import(/* webpackChunkName: "about" */ '../views/logs/Rmi') + }, { path: '/rules/http', name: 'HttpRules', @@ -39,6 +44,11 @@ const routes = [ path: '/rules/mysql', name: 'MysqlRules', component: () => import(/* webpackChunkName: "about" */ '../views/rules/Mysql') + }, + { + path: '/rules/rmi', + name: 'RmiRules', + component: () => import(/* webpackChunkName: "about" */ '../views/rules/Rmi') } ] diff --git a/frontend/src/views/logs/Http.vue b/frontend/src/views/logs/Http.vue index 4326286..74328a9 100644 --- a/frontend/src/views/logs/Http.vue +++ b/frontend/src/views/logs/Http.vue @@ -131,6 +131,7 @@ const columns = [ title: 'PATH', dataIndex: 'path', key: 'path', + ellipsis: true, scopedSlots: { filterDropdown: 'filterDropdown', filterIcon: 'filterIcon', diff --git a/frontend/src/views/logs/Rmi.vue b/frontend/src/views/logs/Rmi.vue new file mode 100644 index 0000000..5e41abc --- /dev/null +++ b/frontend/src/views/logs/Rmi.vue @@ -0,0 +1,171 @@ + + + \ No newline at end of file diff --git a/frontend/src/views/rules/Mysql.vue b/frontend/src/views/rules/Mysql.vue index 29068e8..de79e1b 100644 --- a/frontend/src/views/rules/Mysql.vue +++ b/frontend/src/views/rules/Mysql.vue @@ -137,11 +137,6 @@ - - - {{ files.length > 30 ? files.substr(0, 30) + "..." : files }} - - @@ -221,9 +216,7 @@ const columns = [ title: 'FILES', dataIndex: 'files', key: 'files', - scopedSlots: { - customRender: "files" - } + ellipsis: true }, { title: 'PUSH TO CLIENT', diff --git a/frontend/src/views/rules/Rmi.vue b/frontend/src/views/rules/Rmi.vue new file mode 100644 index 0000000..b22ef5d --- /dev/null +++ b/frontend/src/views/rules/Rmi.vue @@ -0,0 +1,337 @@ + + + \ No newline at end of file diff --git a/pkg/dns/dns.go b/pkg/dns/dns.go index f07051e..54f0a11 100644 --- a/pkg/dns/dns.go +++ b/pkg/dns/dns.go @@ -82,7 +82,7 @@ func (s *Server) Run() { log.Error("DNS record(rule_id:%s) created failed :%s", _rule.Name, err.Error()) return nil, nil } - log.Trace("DNS record(id:%d) has been created", r.ID) + log.Info("DNS record(id:%d,rule:%s,remote_ip:%s) has been created", r.ID, _rule.Name, ip) //only send to client when this connection recorded first time. if _rule.PushToClient { diff --git a/pkg/dns/record.go b/pkg/dns/record.go index fd207a5..c59da51 100644 --- a/pkg/dns/record.go +++ b/pkg/dns/record.go @@ -42,7 +42,7 @@ func newRecord(rule *Rule, flag, domain, remoteIp, ipArea string) (r *Record, er return r, err } -func List(c *gin.Context) { +func ListRecords(c *gin.Context) { var ( dnsRecord Record res []Record @@ -80,6 +80,11 @@ func List(c *gin.Context) { }) return } + + if order != "desc" && order != "asc" { + order = "desc" + } + if err := db.Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil { c.JSON(400, gin.H{ "status": "failed", diff --git a/pkg/dns/rule.go b/pkg/dns/rule.go index bf0fd4b..6cdc614 100644 --- a/pkg/dns/rule.go +++ b/pkg/dns/rule.go @@ -107,6 +107,10 @@ func ListRules(c *gin.Context) { return } + if order != "desc" && order != "asc" { + order = "desc" + } + if err := db.Order("rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil { c.JSON(400, gin.H{ "status": "failed", diff --git a/pkg/mysql/mysql.go b/pkg/mysql/mysql.go index bde3e31..58f4108 100644 --- a/pkg/mysql/mysql.go +++ b/pkg/mysql/mysql.go @@ -125,7 +125,7 @@ func (s *Server) ConnectionClosed(c *vmysql.Conn) { log.Error("MySQL record(rule_id:%s) created failed :%s", _rule.Name, err.Error()) return } - log.Trace("MySQL record(id:%d) has been created", r.ID) + log.Info("MySQL record(id:%d,rule:%s,remote_ip:%s) has been created", r.ID, _rule.Name, ip) //only send to client when this connection recorded first time. if _rule.PushToClient { @@ -267,10 +267,10 @@ func (s *Server) Run() { var authServer = &vmysql.AuthServerNone{} var err error - log.Info("Starting Mysql Server at %s", s.Addr) + log.Info("Starting MySQL Server at %s", s.Addr) s.listener, err = vmysql.NewListener("tcp", s.Addr, authServer, s, s.VersionString, 0, 0) if err != nil { - log.Error("New Mysql Server failed: %s", err) + log.Error("New MySQL Server failed: %s", err) os.Exit(-1) } diff --git a/pkg/mysql/record.go b/pkg/mysql/record.go index 2c18f7d..b70511b 100644 --- a/pkg/mysql/record.go +++ b/pkg/mysql/record.go @@ -49,13 +49,14 @@ func newRecord(rule *Rule, flag, username, clientName, clientOS, remoteIp, ipAre return r, err } -func List(c *gin.Context) { +func ListRecords(c *gin.Context) { var ( mysqlRecord Record res []Record count int64 order = c.Query("order") ) + if err := c.ShouldBind(&mysqlRecord); err != nil { c.JSON(400, gin.H{ "status": "failed", @@ -90,6 +91,11 @@ func List(c *gin.Context) { }) return } + + if order != "desc" && order != "asc" { + order = "desc" + } + if err := db.Preload("Files").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil { c.JSON(400, gin.H{ "status": "failed", diff --git a/pkg/mysql/rule.go b/pkg/mysql/rule.go index 5ca90bf..95eac1d 100644 --- a/pkg/mysql/rule.go +++ b/pkg/mysql/rule.go @@ -89,6 +89,11 @@ func ListRules(c *gin.Context) { }) return } + + if order != "desc" && order != "asc" { + order = "desc" + } + if err := db.Order("rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil { c.JSON(400, gin.H{ "status": "failed", diff --git a/pkg/rhttp/http.go b/pkg/rhttp/http.go index 3ee74c6..a27b711 100644 --- a/pkg/rhttp/http.go +++ b/pkg/rhttp/http.go @@ -157,7 +157,7 @@ func (s *Server) Receive(c *gin.Context) { c.String(code, compileTpl(c, _rule.ResponseBody)) return } - log.Trace("HTTP record(id:%d) has been created", r.ID) + log.Info("HTTP record(id:%d,rule:%s,remote_ip:%s) has been created", r.ID, _rule.Name, ip) //only send to client when this connection recorded first time. if _rule.PushToClient { diff --git a/pkg/rhttp/record.go b/pkg/rhttp/record.go index bc48b03..52e0fda 100644 --- a/pkg/rhttp/record.go +++ b/pkg/rhttp/record.go @@ -89,6 +89,10 @@ func ListRecords(c *gin.Context) { return } + if order != "desc" && order != "asc" { + order = "desc" + } + if err := db.Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil { c.JSON(400, gin.H{ "status": "failed", diff --git a/pkg/rhttp/rule.go b/pkg/rhttp/rule.go index d57d464..cc48ef4 100644 --- a/pkg/rhttp/rule.go +++ b/pkg/rhttp/rule.go @@ -108,6 +108,10 @@ func ListRules(c *gin.Context) { return } + if order != "desc" && order != "asc" { + order = "desc" + } + if err := db.Order("rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil { c.JSON(400, gin.H{ "status": "failed", diff --git a/pkg/rmi/config.go b/pkg/rmi/config.go new file mode 100644 index 0000000..19129fb --- /dev/null +++ b/pkg/rmi/config.go @@ -0,0 +1,6 @@ +package rmi + +type Config struct { + Enable bool + Addr string +} diff --git a/pkg/rmi/record.go b/pkg/rmi/record.go new file mode 100644 index 0000000..20b6661 --- /dev/null +++ b/pkg/rmi/record.go @@ -0,0 +1,103 @@ +package rmi + +import ( + "strconv" + "time" + + "github.com/gin-gonic/gin" + "github.com/li4n0/revsuit/internal/database" + "github.com/li4n0/revsuit/internal/notice" + "github.com/li4n0/revsuit/internal/record" +) + +var _ record.Record = (*Record)(nil) + +type Record struct { + Path string `form:"path" json:"path"` + record.BaseRecord + Rule Rule `gorm:"foreignKey:RuleName;references:Name;constraint:OnUpdate:CASCADE,OnDelete:SET NULL;" form:"-" json:"-" notice:"-"` +} + +func (Record) TableName() string { + return "rmi_records" +} + +func (r Record) Notice() { + notice.Notice(r) +} + +func NewRecord(rule *Rule, flag, path, ip, area string) (r *Record, err error) { + r = &Record{ + BaseRecord: record.BaseRecord{ + Flag: flag, + RemoteIP: ip, + IpArea: area, + RequestTime: time.Now(), + }, + Path: path, + Rule: *rule, + } + err = database.DB.Create(r).Error + return r, err +} + +func ListRecords(c *gin.Context) { + var ( + rmiRecord Record + res []Record + count int64 + order = c.Query("order") + ) + + if err := c.ShouldBind(&rmiRecord); err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err, + "result": nil, + }) + return + } + + db := database.DB.Model(&rmiRecord) + if rmiRecord.Flag != "" { + db.Where("flag = ?", rmiRecord.Flag) + } + if rmiRecord.Path != "" { + db.Where("path like ?", "%"+rmiRecord.Path+"%") + } + if rmiRecord.RemoteIP != "" { + db.Where("remote_ip = ?", rmiRecord.RemoteIP) + } + if rmiRecord.RuleName != "" { + db.Where("rule_name = ?", rmiRecord.RuleName) + } + + page, err := strconv.Atoi(c.Query("page")) + if err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err.Error(), + "result": nil, + }) + return + } + + if order != "desc" && order != "asc" { + order = "desc" + } + + if err := db.Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err.Error(), + "data": nil, + }) + return + } + + c.JSON(200, gin.H{ + "status": "succeed", + "error": nil, + "result": gin.H{"count": count, "data": res}, + }) +} diff --git a/pkg/rmi/rmi.go b/pkg/rmi/rmi.go new file mode 100644 index 0000000..3fbdb58 --- /dev/null +++ b/pkg/rmi/rmi.go @@ -0,0 +1,158 @@ +package rmi + +import ( + "bytes" + "encoding/binary" + "net" + "strconv" + "strings" + "sync" + "time" + + "github.com/li4n0/revsuit/internal/database" + "github.com/li4n0/revsuit/internal/qqwry" + log "unknwon.dev/clog/v2" +) + +type Server struct { + Config + rules []*Rule + rulesLock sync.RWMutex +} + +var ( + server *Server + once sync.Once +) + +func GetServer() *Server { + once.Do(func() { + server = &Server{rulesLock: sync.RWMutex{}} + }) + return server +} + +func (s *Server) getRules() []*Rule { + defer s.rulesLock.RUnlock() + s.rulesLock.RLock() + return s.rules +} + +func (s *Server) updateRules() error { + db := database.DB.Model(new(Rule)) + s.rulesLock.Lock() + db.Order("rank desc").Find(&s.rules) + s.rulesLock.Unlock() + return nil +} + +func (s *Server) handleConnection(conn net.Conn) { + defer conn.Close() + + ip, port, _ := net.SplitHostPort(conn.RemoteAddr().String()) + if err := conn.SetDeadline(time.Now().Add(time.Second * 30)); err != nil { + log.Error("RMI set connection deadline error:%v", err.Error()) + } + + buf := make([]byte, 1024) + _, err := conn.Read(buf) + if err != nil { + log.Error("RMI read connection error:%v", err.Error()) + } + + if !bytes.Contains(buf, []byte{0x4a, 0x52, 0x4d, 0x49}) { + return + } + + send := []byte{0x4e} + bs := make([]byte, 8) + binary.BigEndian.PutUint16(bs, uint16(len(ip))) + send = append(send, bs...) + send = append(send, []byte(ip)...) + send = append(send, []byte{0x00, 0x00}...) + uintPort, _ := strconv.Atoi(port) + bs = make([]byte, 8) + binary.BigEndian.PutUint16(bs, uint16(uintPort)) + send = append(send, bs...) + + _, err = conn.Write(send) + if err != nil { + log.Error("RMI write connection error: %v", err.Error()) + } + + data := make([]byte, 512) + + for length := 0; length < 50; { + n, err := conn.Read(data) + if err != nil { + log.Error("RMI read connection error: %v", err.Error()) + } + length += n + } + + frags := bytes.Split(data, []byte{0xdf, 0x74}) + path := strings.TrimRight(string(frags[len(frags)-1][2:]), "\x00") + + for _, _rule := range s.getRules() { + flag, flagGroup := _rule.Match(path) + if flag == "" { + continue + } + + area := qqwry.Area(ip) + + // create new record + r, err := NewRecord(_rule, flag, path, ip, area) + if err != nil { + log.Error("RMI record(rule_id:%d) created failed :%s", _rule.ID, err.Error()) + return + } + log.Info("RMI record(id:%d,rule:%s,remote_ip:%s) has been created", r.ID, _rule.Name, ip) + + //only send to client when this connection recorded first time. + if _rule.PushToClient { + if flagGroup != "" { + var count int64 + database.DB.Where("rule_name=? and raw like ?", _rule.Name, "%"+flagGroup+"%").Model(&Record{}).Count(&count) + if count <= 1 { + r.PushToClient() + log.Trace("RMI record(id:%d) has been put to client message queue", r.ID) + } + } + r.PushToClient() + log.Trace("RMI record(id:%d) has been put to client message queue", r.ID) + } + + //send notice + if _rule.Notice { + go func() { + r.Notice() + log.Trace("RMI record(id:%d) notice has been sent", r.ID) + }() + } + } +} + +func (s *Server) Run() { + if err := s.updateRules(); err != nil { + log.Fatal(err.Error()) + } + + // run server + log.Info("Starting RMI Server at %v", s.Addr) + + listener, err := net.Listen("tcp", s.Addr) + if err != nil { + log.Fatal(err.Error()) + } + + for { + tcpConn, err := listener.Accept() + if err != nil { + log.Error("RMI accept connection error: %v", err.Error()) + continue + } + go s.handleConnection(tcpConn) + } + +} diff --git a/pkg/rmi/rule.go b/pkg/rmi/rule.go new file mode 100644 index 0000000..e2bf0f9 --- /dev/null +++ b/pkg/rmi/rule.go @@ -0,0 +1,193 @@ +package rmi + +import ( + "strconv" + + "github.com/gin-gonic/gin" + "github.com/li4n0/revsuit/internal/database" + "github.com/li4n0/revsuit/internal/rule" + "gorm.io/gorm/clause" + log "unknwon.dev/clog/v2" +) + +// Http rule struct +type Rule struct { + rule.BaseRule +} + +func (Rule) TableName() string { + return "rmi_rules" +} + +// New rmi rule struct +func NewRule(name, flagFormat string, pushToClient, notice bool) *Rule { + return &Rule{ + BaseRule: rule.BaseRule{ + Name: name, + FlagFormat: flagFormat, + PushToClient: pushToClient, + Notice: notice, + }, + } +} + +// Create or update the rmi rule in database and ruleSet +func (r *Rule) CreateOrUpdate() (err error) { + db := database.DB.Model(r) + err = db.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "id"}}, + DoUpdates: clause.AssignmentColumns( + []string{ + "name", + "flag_format", + "rank", + "push_to_client", + "notice", + }), + }).Create(r).Error + if err != nil { + return + } + + err = GetServer().updateRules() + return err +} + +// Delete the rmi rule in database and ruleSet +func (r *Rule) Delete() (err error) { + db := database.DB.Model(r) + err = db.Delete(r).Error + if err != nil { + return + } + + err = GetServer().updateRules() + return err +} + +// List all rmi rules those satisfy the filter +func ListRules(c *gin.Context) { + var ( + rmiRule Rule + res []Rule + count int64 + order = c.Query("order") + ) + + if err := c.ShouldBind(&rmiRule); err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err, + "result": nil, + }) + return + } + + db := database.DB.Model(&rmiRule) + if rmiRule.Name != "" { + db.Where("name = ?", rmiRule.Name) + } + db.Count(&count) + + page, err := strconv.Atoi(c.Query("page")) + if err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err.Error(), + "result": nil, + }) + return + } + + if order != "desc" && order != "asc" { + order = "desc" + } + + if err := db.Order("rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * 10).Limit(10).Find(&res).Error; err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err.Error(), + "data": nil, + }) + return + } + + c.JSON(200, gin.H{ + "status": "succeed", + "error": nil, + "result": gin.H{"count": count, "data": res}, + }) +} + +// Create or update rmi rule from user submit +func UpsertRules(c *gin.Context) { + var ( + rmiRule Rule + update bool + ) + + if err := c.ShouldBind(&rmiRule); err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err.Error(), + "data": nil, + }) + return + } + + if rmiRule.ID != 0 { + update = true + } + + if err := rmiRule.CreateOrUpdate(); err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err.Error(), + "result": nil, + }) + return + } + + if update { + log.Trace("RMI rule(id:%d) has been updated", rmiRule.ID) + } else { + log.Trace("RMI rule(id:%d) has been created", rmiRule.ID) + } + + c.JSON(200, gin.H{ + "status": "succeed", + "error": nil, + "result": nil, + }) +} + +// Delete rmi rule from user submit +func DeleteRules(c *gin.Context) { + var rmiRule Rule + + if err := c.ShouldBind(&rmiRule); err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err.Error(), + "data": nil, + }) + return + } + + if err := rmiRule.Delete(); err != nil { + c.JSON(400, gin.H{ + "status": "failed", + "error": err.Error(), + "data": nil, + }) + return + } + + log.Trace("RMI rule(id:%d) has been deleted", rmiRule.ID) + + c.JSON(200, gin.H{ + "status": "succeed", + "error": nil, + "data": nil, + }) +} diff --git a/pkg/server/config.go b/pkg/server/config.go index df44448..7571a24 100644 --- a/pkg/server/config.go +++ b/pkg/server/config.go @@ -4,6 +4,7 @@ import ( "github.com/li4n0/revsuit/pkg/dns" "github.com/li4n0/revsuit/pkg/mysql" "github.com/li4n0/revsuit/pkg/rhttp" + "github.com/li4n0/revsuit/pkg/rmi" ) type noticeConfig struct { @@ -21,5 +22,6 @@ type Config struct { Notice noticeConfig rhttp.Config DNS dns.Config - Mysql mysql.Config + MySQL mysql.Config + RMI rmi.Config } diff --git a/pkg/server/router.go b/pkg/server/router.go index f6c858d..6247fcf 100644 --- a/pkg/server/router.go +++ b/pkg/server/router.go @@ -9,6 +9,7 @@ import ( "github.com/li4n0/revsuit/pkg/dns" "github.com/li4n0/revsuit/pkg/mysql" "github.com/li4n0/revsuit/pkg/rhttp" + "github.com/li4n0/revsuit/pkg/rmi" log "unknwon.dev/clog/v2" ) @@ -52,10 +53,13 @@ func (revsuit *Revsuit) registerHttpRouter() { httpGroup.GET("", rhttp.ListRecords) dnsGroup := recordGroup.Group("/dns") - dnsGroup.GET("", dns.List) + dnsGroup.GET("", dns.ListRecords) mysqlGroup := recordGroup.Group("/mysql") - mysqlGroup.GET("", mysql.List) + mysqlGroup.GET("", mysql.ListRecords) + + rmiGroup := recordGroup.Group("/rmi") + rmiGroup.GET("", rmi.ListRecords) // init rule router group ruleGroup := revsuit.http.ApiGroup.Group("/rule") @@ -75,6 +79,11 @@ func (revsuit *Revsuit) registerHttpRouter() { mysqlGroup.POST("", mysql.UpsertRules) mysqlGroup.DELETE("", mysql.DeleteRules) + rmiGroup = ruleGroup.Group("/rmi") + rmiGroup.GET("", rmi.ListRules) + rmiGroup.POST("", rmi.UpsertRules) + rmiGroup.DELETE("", rmi.DeleteRules) + // init file router group fileGroup := revsuit.http.ApiGroup.Group("/file") fileGroup.GET("/mysql/:id", mysql.GetFile) diff --git a/pkg/server/server.go b/pkg/server/server.go index d2ea5a7..025b683 100644 --- a/pkg/server/server.go +++ b/pkg/server/server.go @@ -7,6 +7,7 @@ import ( "github.com/li4n0/revsuit/pkg/dns" "github.com/li4n0/revsuit/pkg/mysql" http "github.com/li4n0/revsuit/pkg/rhttp" + "github.com/li4n0/revsuit/pkg/rmi" "gorm.io/gorm/logger" log "unknwon.dev/clog/v2" ) @@ -15,6 +16,7 @@ type Revsuit struct { http *http.Server dns *dns.Server mysql *mysql.Server + rmi *rmi.Server } func initDatabase(dsn string) { @@ -51,6 +53,14 @@ func initDatabase(dsn string) { if err != nil { log.Fatal(err.Error()) } + err = database.DB.AutoMigrate(&rmi.Record{}) + if err != nil { + log.Fatal(err.Error()) + } + err = database.DB.AutoMigrate(&rmi.Rule{}) + if err != nil { + log.Fatal(err.Error()) + } } @@ -122,9 +132,13 @@ func New(c *Config) *Revsuit { if c.DNS.Enable { s.dns = dns.GetServer() } - if c.Mysql.Enable { + if c.MySQL.Enable { s.mysql = mysql.GetServer() - s.mysql.Config = c.Mysql + s.mysql.Config = c.MySQL + } + if c.RMI.Enable { + s.rmi = rmi.GetServer() + s.rmi.Config = c.RMI } if c.Addr != "" { @@ -149,6 +163,9 @@ func (revsuit *Revsuit) Run() { if revsuit.mysql != nil { go revsuit.mysql.Run() } + if revsuit.rmi != nil { + go revsuit.rmi.Run() + } revsuit.http.Run() }