archive/revsuit
1
0
Fork 0
mirror of https://github.com/Li4n0/revsuit.git synced 2026-09-21 22:30:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(database): support postgres (#43)

Browse Source 
Co-authored-by: Li4n0 <[email protected]>
This commit is contained in:
becivells
2021-12-24 22:05:12 +08:00
committed by GitHub
co-authored by Li4n0
parent 3242a49b27
commit 8b30f83075
18 changed files with 224 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/dns/dns.go
+1 -1
View File
@@ -57,7 +57,7 @@ func (s *Server) UpdateRules() error {
db := database.DB.Model(new(Rule))
defer s.rulesLock.Unlock()
s.rulesLock.Lock()
return errors.Wrap(db.Order("`rank` desc").Find(&s.rules).Error, "DNS update rules error")
return errors.Wrap(db.Order("base_rank desc").Find(&s.rules).Error, "DNS update rules error")
}
func newSet(_rule *Rule, name, value, ip string, _type newdns.Type) []newdns.Set {
pkg/dns/rule.go
+2 -2
View File
@@ -47,7 +47,7 @@ func (r *Rule) CreateOrUpdate() (err error) {
[]string{
"name",
"flag_format",
"rank",
"base_rank",
"type",
"value",
"ttl",
@@ -120,7 +120,7 @@ func ListRules(c *gin.Context) {
order = "desc"
}
if err := db.Order("`rank` desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
if err := db.Order("base_rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
c.JSON(400, gin.H{
"status": "failed",
"error": err.Error(),
pkg/ftp/ftp.go
+1 -1
View File
@@ -71,7 +71,7 @@ func (s *Server) UpdateRules() error {
db := database.DB.Model(new(Rule))
defer s.rulesLock.Unlock()
s.rulesLock.Lock()
return errors.Wrap(db.Order("`rank` desc").Find(&s.rules).Error, "FTP update rules error")
return errors.Wrap(db.Order("base_rank desc").Find(&s.rules).Error, "FTP update rules error")
}
func getClientPasvConnAddress(ip, port string) string {
pkg/ftp/rule.go
+2 -2
View File
@@ -43,7 +43,7 @@ func (r *Rule) CreateOrUpdate() (err error) {
[]string{
"name",
"flag_format",
"rank",
"base_rank",
"pasv_address",
"data",
"push_to_client",
@@ -115,7 +115,7 @@ func ListRules(c *gin.Context) {
order = "desc"
}
if err := db.Order("`rank` desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
if err := db.Order("base_rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
c.JSON(400, gin.H{
"status": "failed",
"error": err.Error(),
pkg/mysql/mysql.go
+1 -1
View File
@@ -52,7 +52,7 @@ func (s *Server) UpdateRules() error {
db := database.DB.Model(new(Rule))
defer s.rulesLock.Unlock()
s.rulesLock.Lock()
return errors.Wrap(db.Order("`rank` desc").Find(&s.rules).Error, "MySQL update rules error")
return errors.Wrap(db.Order("base_rank desc").Find(&s.rules).Error, "MySQL update rules error")
}
// NewConnection is part of the mysql.Handler interface.
pkg/mysql/rule.go
+2 -2
View File
@@ -30,7 +30,7 @@ func (r *Rule) CreateOrUpdate() (err error) {
[]string{
"name",
"flag_format",
"rank",
"base_rank",
"files",
"exploit_jdbc_client",
"payloads",
@@ -103,7 +103,7 @@ func ListRules(c *gin.Context) {
order = "desc"
}
if err := db.Order("`rank` desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
if err := db.Order("base_rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
c.JSON(400, gin.H{
"status": "failed",
"error": err.Error(),
pkg/rhttp/http.go
+1 -1
View File
@@ -87,7 +87,7 @@ func (s *Server) UpdateRules() error {
db := database.DB.Model(new(Rule))
defer s.rulesLock.Unlock()
s.rulesLock.Lock()
return db.Order("`rank` desc").Find(&s.rules).Error
return db.Order("base_rank desc").Find(&s.rules).Error
}
func (s *Server) startHttpServer() {
pkg/rhttp/rule.go
+2 -2
View File
@@ -46,7 +46,7 @@ func (r *Rule) CreateOrUpdate() (err error) {
[]string{
"name",
"flag_format",
"rank",
"base_rank",
"response_status_code",
"response_headers",
"response_body",
@@ -121,7 +121,7 @@ func ListRules(c *gin.Context) {
order = "desc"
}
if err := db.Order("`rank` desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
if err := db.Order("base_rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
c.JSON(400, gin.H{
"status": "failed",
"error": err.Error(),
pkg/rmi/rmi.go
+1 -1
View File
@@ -45,7 +45,7 @@ func (s *Server) UpdateRules() error {
db := database.DB.Model(new(Rule))
defer s.rulesLock.Unlock()
s.rulesLock.Lock()
return errors.Wrap(db.Order("`rank` desc").Find(&s.rules).Error, "RMI update rules error")
return errors.Wrap(db.Order("base_rank desc").Find(&s.rules).Error, "RMI update rules error")
}
func (s *Server) handleConnection(conn net.Conn) {
pkg/rmi/rule.go
+2 -2
View File
@@ -40,7 +40,7 @@ func (r *Rule) CreateOrUpdate() (err error) {
[]string{
"name",
"flag_format",
"rank",
"base_rank",
"push_to_client",
"notice",
}),
@@ -112,7 +112,7 @@ func ListRules(c *gin.Context) {
order = "desc"
}
if err := db.Order("`rank` desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
if err := db.Order("base_rank desc").Order("id" + " " + order).Count(&count).Offset((page - 1) * pageSize).Limit(pageSize).Find(&res).Error; err != nil {
c.JSON(400, gin.H{
"status": "failed",
"error": err.Error(),
pkg/server/server.go
+2 -11
View File
@@ -1,11 +1,8 @@
package server
import (
"strings"
"sync"
"github.com/pkg/errors"
"github.com/gin-gonic/gin"
"github.com/li4n0/revsuit/internal/database"
"github.com/li4n0/revsuit/internal/file"
@@ -63,14 +60,8 @@ func initDatabase(dsn string) {
Level: log.LevelInfo,
})
var err error
if strings.Contains(dsn, "@tcp") {
err = database.InitDB(database.Mysql, dsn)
} else if strings.Contains(dsn, ".db") {
err = database.InitDB(database.Sqlite, dsn)
} else {
err = errors.New("unsupported database")
}
err := database.InitDB(dsn)
if err != nil {
log.Fatal(err.Error())
}