sacrebleu-dns/utils/sql.go

102 lines
2.9 KiB
Go
Raw Normal View History

2020-12-13 04:01:04 +01:00
package utils
import (
"fmt"
"time"
2020-12-28 01:34:26 +01:00
"github.com/outout14/sacrebleu-api/api/types"
2020-12-13 04:01:04 +01:00
"github.com/sirupsen/logrus"
2020-12-23 01:09:10 +01:00
"gorm.io/driver/mysql"
"gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/logger"
2020-12-13 04:01:04 +01:00
)
//DB SQL database as global var
2020-12-23 01:09:10 +01:00
var db *gorm.DB
2020-12-13 04:01:04 +01:00
2020-12-28 00:05:16 +01:00
//SQLDatabase Initialize the SQL Database
2020-12-14 23:20:24 +01:00
//Requires a conf struct
2020-12-28 00:05:16 +01:00
func SQLDatabase(conf *Conf) *gorm.DB {
logrus.WithFields(logrus.Fields{"database": conf.Database.Db, "driver": conf.Database.Type}).Infof("SQL : Connection to DB")
var gormLogLevel logger.LogLevel
//Set GORM log level based on conf AppMode
if conf.AppMode != "production" {
gormLogLevel = logger.Info
} else {
gormLogLevel = logger.Silent
}
2020-12-28 00:05:16 +01:00
//Connect to the Database
if conf.Database.Type == "postgresql" {
dsn := fmt.Sprintf("user=%s password=%s host=%s port=%s database=%s sslmode=disable", conf.Database.Username, conf.Database.Password, conf.Database.Host, conf.Database.Port, conf.Database.Db)
2020-12-28 00:05:16 +01:00
DB, err := gorm.Open(postgres.Open(dsn), &gorm.Config{
Logger: logger.Default.LogMode(gormLogLevel),
})
2020-12-28 00:05:16 +01:00
CheckErr(err)
db = DB
return DB
}
2020-12-28 00:05:16 +01:00
dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?charset=utf8mb4&parseTime=True&loc=Local", conf.Database.Username, conf.Database.Password, conf.Database.Host, conf.Database.Port, conf.Database.Db)
2020-12-13 04:01:04 +01:00
2020-12-28 00:05:16 +01:00
DB, err := gorm.Open(mysql.Open(dsn), &gorm.Config{
Logger: logger.Default.LogMode(gormLogLevel),
})
2020-12-23 01:09:10 +01:00
CheckErr(err)
2020-12-28 00:05:16 +01:00
db = DB
return DB
2020-12-13 04:01:04 +01:00
}
2020-12-14 23:20:24 +01:00
//Check for a record in the SQL database
2020-12-28 01:34:26 +01:00
func sqlCheckForRecord(redisKey string, dKey string, entry types.Record) ([]types.Record, bool) {
var records []types.Record
2020-12-13 04:01:04 +01:00
2020-12-28 01:34:26 +01:00
rows, err := db.Where("fqdn = ? AND type = ?", dKey, entry.Qtype).Model(&types.Record{}).Rows()
2020-12-28 00:05:16 +01:00
if err != nil {
return records, true
2020-12-13 04:01:04 +01:00
}
2020-12-28 00:05:16 +01:00
defer rows.Close()
for rows.Next() {
2020-12-28 01:34:26 +01:00
var entry types.Record
2020-12-28 00:05:16 +01:00
db.ScanRows(rows, &entry)
2020-12-28 00:05:16 +01:00
if entry.Content != "" { //If Record content not empty
records = append(records, entry)
}
}
//Cache the request in Redis if any result
_ = redisSet(redisDb, redisKey, 30*time.Second, records) //Set it in the Redis database for 30sec
return records, false
2020-12-13 04:01:04 +01:00
}
2020-12-14 23:20:24 +01:00
//Check for a wildcard record in the SQL database
2020-12-28 01:34:26 +01:00
func sqlCheckForReverse6Wildcard(redisKey string, dKey string, entry types.Record) []types.Record {
2020-12-13 04:01:04 +01:00
returnedEntry := entry
2020-12-23 02:01:18 +01:00
rows, err := db.Table("records").Select("id", "content", "fqdn").Where("fqdn LIKE ?", "*%.ip6.arpa.").Rows()
2020-12-23 01:09:10 +01:00
DbgErr(err) //Check for empty row or non important error
2020-12-13 04:01:04 +01:00
2020-12-28 01:34:26 +01:00
var records []types.Record
2020-12-28 00:05:16 +01:00
2020-12-14 23:20:24 +01:00
//For each result check if it match the reverse IP
2020-12-23 01:09:10 +01:00
for rows.Next() {
2020-12-28 00:05:16 +01:00
rows.Scan(&returnedEntry.ID, &returnedEntry.Content, &returnedEntry.Fqdn)
2020-12-13 04:01:04 +01:00
CheckErr(err)
2020-12-14 23:20:24 +01:00
//Check if the record is matching the reversed IP
2020-12-13 04:01:04 +01:00
if checkReverse6(entry, returnedEntry) {
logrus.Debug("REVERSE : Correct wildcard reverse.")
//Cache the request in Redis if any result
_ = redisSet(redisDb, redisKey, 10*time.Second, returnedEntry)
2020-12-28 00:05:16 +01:00
records = append(records, entry)
2020-12-13 04:01:04 +01:00
}
logrus.Debug("REVERSE : WRONG wildcard reverse .")
2020-12-13 04:01:04 +01:00
}
2020-12-28 00:05:16 +01:00
return records
2020-12-13 04:01:04 +01:00
}