From 2a2199e3eb3f9383f4fe98c0cac20173956821df Mon Sep 17 00:00:00 2001 From: Paul Lecuq Date: Sun, 16 Oct 2022 23:27:57 +0200 Subject: [PATCH] updated models --- src/database/main.go | 19 ++-- src/models/as.go | 56 ++++++++++ src/models/cfg.go | 30 ++--- src/models/city.go | 54 +++++++++ src/models/country.go | 52 +++++++++ src/models/event.go | 57 ++++++++++ src/models/hostinfo.go | 70 ++++++++---- src/models/ip.go | 232 ++++++++++++++++++++++++--------------- src/models/models.go | 65 +++++++++++ src/models/scanresult.go | 65 +++++++++++ src/routers/funcs.go | 44 ++++++-- src/routers/main.go | 3 + src/zmqrouter/main.go | 24 ++-- 13 files changed, 614 insertions(+), 157 deletions(-) create mode 100644 src/models/as.go create mode 100644 src/models/city.go create mode 100644 src/models/country.go create mode 100644 src/models/event.go create mode 100644 src/models/models.go create mode 100644 src/models/scanresult.go diff --git a/src/database/main.go b/src/database/main.go index f1c876d..12023c7 100644 --- a/src/database/main.go +++ b/src/database/main.go @@ -10,17 +10,14 @@ import ( "git.paulbsd.com/paulbsd/ipbl/src/models" _ "github.com/lib/pq" "xorm.io/xorm" + "xorm.io/xorm/caches" + "xorm.io/xorm/dialects" "xorm.io/xorm/names" ) func Initialize(ctx *context.Context, cfg *config.Config) (err error) { var databaseEngine = "postgres" - var tables = []interface{}{models.IP{}, - models.Cfg{}, - models.CfgSet{}, - models.CfgTrustlist{}, - models.CfgZMQ{}, - models.HostInfo{}} + cacher := caches.NewLRUCacher(caches.NewMemoryStore(), 1000) cfg.Db, err = xorm.NewEngine(databaseEngine, fmt.Sprintf("%s://%s:%s@%s/%s", @@ -33,18 +30,18 @@ func Initialize(ctx *context.Context, cfg *config.Config) (err error) { log.Fatalln(err) } cfg.Db.SetMapper(names.GonicMapper{}) + cfg.Db.SetQuotePolicy(dialects.QuotePolicyReserved) cfg.Db.ShowSQL(cfg.Switchs.Debug) + cfg.Db.SetDefaultCacher(cacher) + if cfg.Switchs.Drop { log.Println("Dropping tables") - cfg.Db.DropTables(tables) + models.DropTables(ctx, cfg) os.Exit(0) } log.Println("Syncing tables") - for _, table := range tables { - cfg.Db.CreateTables(table) - err = cfg.Db.Sync2(table) - } + err = models.NewEngine(ctx, cfg) return } diff --git a/src/models/as.go b/src/models/as.go new file mode 100644 index 0000000..041345f --- /dev/null +++ b/src/models/as.go @@ -0,0 +1,56 @@ +package models + +import "xorm.io/xorm" + +func (as *AutonomousSystem) GetOrCreate(session *xorm.Session) (apias *APIAutonomousSystem, err error) { + has, err := session.Get(as) + if err != nil { + return + } + if !has { + session.Insert(as) + } else { + session.ID(as.ID).Update(as) + } + as.Get(session) + apias = as.APIFormat() + return +} + +func (as *AutonomousSystem) Get(session *xorm.Session) (apias *APIAutonomousSystem, err error) { + has, err := session.Get(as) + if !has || err != nil { + return + } + apias = as.APIFormat() + return +} + +func (as *AutonomousSystem) APIFormat() *APIAutonomousSystem { + if as == nil { + return &APIAutonomousSystem{} + } + return &APIAutonomousSystem{ + ASID: as.ASID, + ASName: as.ASName, + } +} + +func (as *AutonomousSystem) APIParse(apias APIAutonomousSystem) (err error) { + *as = AutonomousSystem{ + ASID: apias.ASID, + ASName: apias.ASName} + return +} + +type AutonomousSystem struct { + ID int `xorm:"pk autoincr"` + ASID int `xorm:"integer unique(asindex) as_id"` + ASName string `xorm:"text unique(asindex) as_name"` +} + +type APIAutonomousSystem struct { + ID int `json:"-"` + ASID int `json:"as_id"` + ASName string `json:"as_name"` +} diff --git a/src/models/cfg.go b/src/models/cfg.go index 307d9d4..f00253b 100644 --- a/src/models/cfg.go +++ b/src/models/cfg.go @@ -118,17 +118,10 @@ func DiscoverURLS(cfg config.Config, routes []*echo.Route) (Discovery, error) { return disc, nil } -type CfgTrustlist struct { - ID int `xorm:"pk autoincr" json:"-"` - IP string `xorm:"text notnull" json:"ip"` -} - -type CfgZMQ struct { - ID int `xorm:"pk autoincr" json:"-"` - Type string `xorm:"text notnull" json:"type"` - Hostname string `xorm:"text notnull" json:"hostname"` - Port int `json:"port"` - Subscription string `json:"subscription"` +type Cfg struct { + ID int `xorm:"pk autoincr" json:"-"` + Key string `xorm:"text notnull unique" json:"key"` + Value string `xorm:"text default" json:"value"` } type CfgSet struct { @@ -142,10 +135,17 @@ type CfgSet struct { Enabled bool `xorm:"notnull default true" json:"-"` } -type Cfg struct { - ID int `xorm:"pk autoincr" json:"-"` - Key string `xorm:"text notnull unique" json:"key"` - Value string `xorm:"text default" json:"value"` +type CfgTrustlist struct { + ID int `xorm:"pk autoincr" json:"-"` + IP string `xorm:"text notnull" json:"ip"` +} + +type CfgZMQ struct { + ID int `xorm:"pk autoincr" json:"-"` + Type string `xorm:"text notnull" json:"type"` + Hostname string `xorm:"text notnull" json:"hostname"` + Port int `json:"port"` + Subscription string `json:"subscription"` } type Discovery struct { diff --git a/src/models/city.go b/src/models/city.go new file mode 100644 index 0000000..526e41e --- /dev/null +++ b/src/models/city.go @@ -0,0 +1,54 @@ +package models + +import ( + "xorm.io/xorm" +) + +func (city *City) GetOrCreate(session *xorm.Session) (apicity *APICity, err error) { + has, err := session.Get(city) + if err != nil { + return + } + if !has { + session.Insert(city) + } else { + session.ID(city.ID).Update(city) + } + city.Get(session) + apicity = city.APIFormat() + return +} + +func (city *City) Get(session *xorm.Session) (apicity *APICity, err error) { + has, err := session.Get(city) + if !has || err != nil { + return + } + apicity = city.APIFormat() + return +} + +func (city *City) APIFormat() *APICity { + if city == nil { + return &APICity{} + } + return &APICity{ + CityName: city.CityName, + } +} + +func (city *City) APIParse(apicity APICity) (err error) { + *city = City{ + CityName: apicity.CityName} + return +} + +type City struct { + ID int `xorm:"pk autoincr"` + CityName string `xorm:"text unique(cityindex) city_name" json:"city_name"` +} + +type APICity struct { + ID int `json:"-"` + CityName string `json:"city_name"` +} diff --git a/src/models/country.go b/src/models/country.go new file mode 100644 index 0000000..4333f7f --- /dev/null +++ b/src/models/country.go @@ -0,0 +1,52 @@ +package models + +import "xorm.io/xorm" + +func (country *Country) GetOrCreate(session *xorm.Session) (apicountry *APICountry, err error) { + has, err := session.Get(country) + if err != nil { + return + } + if !has { + session.Insert(country) + } else { + session.ID(country.ID).Update(country) + } + apicountry = country.APIFormat() + country.Get(session) + return +} + +func (country *Country) Get(session *xorm.Session) (apicountry *APICountry, err error) { + has, err := session.Get(country) + if !has || err != nil { + return + } + apicountry = country.APIFormat() + return +} + +func (country *Country) APIFormat() *APICountry { + if country == nil { + return &APICountry{} + } + return &APICountry{ + CountryName: country.CountryName, + } +} + +func (country *Country) APIParse(apicountry APICountry) (err error) { + *country = Country{ + CountryName: apicountry.CountryName} + return +} + +type Country struct { + ID int `xorm:"pk autoincr"` + CountryName string `xorm:"text unique(countryindex) country_name" json:"country_name"` +} + +type APICountry struct { + ID int `json:"-"` + CountryName string `json:"country_name"` +} diff --git a/src/models/event.go b/src/models/event.go new file mode 100644 index 0000000..c1f671f --- /dev/null +++ b/src/models/event.go @@ -0,0 +1,57 @@ +package models + +import ( + "database/sql" + "time" + + "git.paulbsd.com/paulbsd/ipbl/src/config" + "xorm.io/xorm" +) + +func (event *Event) Insert(cfg *config.Config) (err error) { + session := cfg.Db.NewSession() + defer session.Close() + _, err = session.Insert(event) + err = session.Commit() + return +} + +func (event *Event) APIFormat() *APIEvent { + if event == nil { + return nil + } + return &APIEvent{ + IP: event.IP.IP, + Src: event.Src, + Hostname: event.Hostname.String, + } +} + +func (event *Event) APIParse(session *xorm.Session, apievent APIEvent) (err error) { + *event = Event{ + IP: &IP{IP: apievent.IP}, + Src: apievent.Src, + Hostname: sql.NullString{ + String: apievent.Hostname, + Valid: true}, + } + event.IP.GetOrCreate(session) + return +} + +type Event struct { + ID int `xorm:"pk autoincr"` + Src string `xorm:"text notnull"` + Hostname sql.NullString `xorm:"text default NULL"` + Mode sql.NullString `xorm:"text default NULL"` + IP *IP `xorm:"int ip_id"` + Created time.Time `xorm:"created notnull"` +} + +type APIEvent struct { + IP string `json:"ip"` + Src string `json:"src"` + Hostname string `json:"hostname"` + Mode string `json:"mode"` + Created string `json:"created"` +} diff --git a/src/models/hostinfo.go b/src/models/hostinfo.go index fd6a823..2f5bcc9 100644 --- a/src/models/hostinfo.go +++ b/src/models/hostinfo.go @@ -1,28 +1,56 @@ package models import ( - "time" + "context" + "database/sql" + + "git.paulbsd.com/paulbsd/ipbl/src/config" + "xorm.io/xorm" ) +func GetIPsWithoutHostInfo(ctx *context.Context, config *config.Config) (apiips []string, err error) { + res, err := config.Db.Query(` + SELECT ip + FROM ip + WHERE ip NOT IN ( + SELECT ip FROM host_info) + ORDER BY RANDOM() + LIMIT 10;`) + for _, r := range res { + apiips = append(apiips, string(r["ip"])) + } + return +} + +func ProcessHostInfo(cfg *config.Config, hostinfos []HostInfo) (err error) { + session := cfg.Db.NewSession() + defer session.Close() + for _, hi := range hostinfos { + var i IP = HostInfoToIP(session, &hi) + i.GetOrCreate(session) + } + session.Commit() + return +} + +func HostInfoToIP(session *xorm.Session, hi *HostInfo) (ip IP) { + ip = IP{ + IP: hi.IP, + Whois: hi.Whois, + Rdns: sql.NullString{ + String: hi.Rdns, + Valid: true}, + City: &City{CityName: hi.City}, + Country: &Country{CountryName: hi.Country}, + AutonomousSystem: &AutonomousSystem{ID: 0}, + } + return +} + type HostInfo struct { - ID int `xorm:"pk autoincr"` - IP string `xorm:"text unique" json:"ip"` - RawWhois string `xorm:"text" json:"raw_whois"` - ReverseDNS string `xorm:"text" json:"reverse_dns"` - City string `xorm:"text" json:"city"` - Country string `xorm:"text" json:"country"` - Created time.Time `xorm:"created notnull"` - Updated time.Time `xorm:"updated notnull"` -} - -type ScanResult struct { - Protocol string `json:"proto"` - PortID int `json:"port_id"` - State string `json:"state"` - ServiceName string `json:"service"` -} - -type ASN struct { - ID int `json:"id"` - Name string `json:"name"` + IP string `json:"ip"` + Whois string `json:"whois"` + Rdns string `json:"rdns"` + City string `json:"city"` + Country string `json:"country"` } diff --git a/src/models/ip.go b/src/models/ip.go index 6f9524d..da235e9 100644 --- a/src/models/ip.go +++ b/src/models/ip.go @@ -3,15 +3,18 @@ package models import ( "context" "database/sql" + "encoding/json" "fmt" + "io" "log" - "net" + "net/http" "time" "git.paulbsd.com/paulbsd/ipbl/src/config" + "xorm.io/xorm" ) -const ScanLimit = 500 +const ScanLimit = 10 var lastday = time.Now().Add(-(time.Hour * 24)) @@ -33,48 +36,43 @@ func GetIPsLast(ctx *context.Context, config *config.Config, interval string) (a return } -func GetIPsWithoutHostInfo(ctx *context.Context, config *config.Config) (apiips []string, err error) { - var ips []IP - err = config.Db.In("ip").GroupBy("ip").Find(&ips) - for _, ml := range ips { - apiips = append(apiips, ml.IP) +func (ip *IP) GetOrCreate(session *xorm.Session) (apiip *APIIP, err error) { + var tmpip = IP{IP: ip.IP} + has, err := session.Get(&tmpip) + ip.City.GetOrCreate(session) + ip.Country.GetOrCreate(session) + ip.AutonomousSystem.GetOrCreate(session) + + if !has { + session.Insert(ip) + } else { + ip.ID = tmpip.ID + session.ID(ip.ID).AllCols().Update(ip) + session.ID(ip.ID).Cols("city", "country", "autonomous_system").Update(ip) } + ip.Get(session) + apiip = ip.APIFormat() return } -func GetIP(ctx *context.Context, cfg *config.Config, ipquery interface{}) (apiip *APIIP, err error) { - session := cfg.Db.NewSession() - defer session.Close() - var ip IP - has, err := session.Where("ip = ?", ipquery).Get(&ip) - if !has { +func (ip *IP) Get(session *xorm.Session) (apiip *APIIP, err error) { + has, err := session.Get(ip) + if !has || err != nil { err = fmt.Errorf("not found") return nil, err } - if err != nil { - return - } apiip = ip.APIFormat() - session.Commit() return } -func (ip *IP) UpdateRDNS() (result string, err error) { - res, err := net.LookupAddr(ip.IP) - if err != nil { - result = "" - return +func (ip *IP) InsertOrUpdate(session *xorm.Session) (numinsert int64, numupdate int64, err error) { + //var ips = []IP{} + //err = session.Where("ip = ?", ip.IP).And("src = ?", ip.Src).And("hostname = ?", ip.Hostname).Find(&ips) + has, err := session.Get(ip) + if has { + session.ID(ip.ID).Update(&IP{}) } - result = res[0] - return -} - -func (ip *IP) InsertOrUpdate(cfg *config.Config) (numinsert int64, numupdate int64, err error) { - session := cfg.Db.NewSession() - defer session.Close() - var ips = []IP{} - err = session.Where("ip = ?", ip.IP).And("src = ?", ip.Src).And("hostname = ?", ip.Hostname).Find(&ips) - if len(ips) > 0 { + /*if len(ips) > 0 { numupdate, err = session.Where("id = ?", ips[0].ID).Cols("updated").Update(&IP{}) if err != nil { log.Println(err) @@ -84,16 +82,16 @@ func (ip *IP) InsertOrUpdate(cfg *config.Config) (numinsert int64, numupdate int if err != nil { log.Println(err) } - } + }*/ session.Commit() return } -func InsertIPBulk(cfg *config.Config, ips *[]IP) (numinsert int64, numupdate int64, err error) { +func InsertIPBulk(session *xorm.Session, ips *[]IP) (numinsert int64, numupdate int64, err error) { for _, ip := range *ips { - numinsert, numupdate, err = ip.InsertOrUpdate(cfg) + numinsert, numupdate, err = ip.InsertOrUpdate(session) } - Cleanup(cfg) + Cleanup(session) return /*var iplist []string hostname := (*ips)[0].Hostname.String @@ -119,20 +117,45 @@ func InsertIPBulk(cfg *config.Config, ips *[]IP) (numinsert int64, numupdate int } func ScanIP(cfg *config.Config) (err error) { + cols := []string{"rdns", "autonomous_system", "city", "country", "updated"} + session := cfg.Db.NewSession() + defer session.Close() + for { orphans := []IP{} - session := cfg.Db.NewSession() - err := session.Begin() - if err != nil { - log.Println(err) - } - if session.Where("rdns IS NULL").Asc("ip").Limit(ScanLimit).Find(&orphans); len(orphans) > 0 { + if session.Asc("updated").Limit(ScanLimit).Find(&orphans); len(orphans) > 0 { + err = session.Begin() + if err != nil { + log.Println(err) + } for _, orphan := range orphans { - reverse, _ := orphan.UpdateRDNS() - log.Printf("%s -> \"%s\"\n", orphan.IP, reverse) - orphan.Rdns.String = reverse + query, err := QueryInfo(orphan.IP) + if err != nil { + log.Println(err) + continue + } + as := &AutonomousSystem{ASID: query.APIAS.ASID, ASName: query.APIAS.ASName} + as.GetOrCreate(session) + orphan.AutonomousSystem = as + + if query.APICity != "" { + var city = &City{CityName: query.APICity} + city.GetOrCreate(session) + orphan.City = city + + } + + if query.APICountry != "" { + var country = &Country{CountryName: query.APICountry} + country.GetOrCreate(session) + orphan.Country = country + } + + orphan.Rdns.String = query.Rdns orphan.Rdns.Valid = true - _, err = session.ID(orphan.ID).Cols("rdns").Update(&orphan) + log.Printf("%s -> \"%s\"\n", orphan.IP, query.Rdns) + + _, err = session.ID(orphan.ID).Cols(cols...).Update(&orphan) if err != nil { session.Rollback() @@ -146,16 +169,42 @@ func ScanIP(cfg *config.Config) (err error) { } else { time.Sleep(10 * time.Second) } - session.Close() + } } -func Cleanup(cfg *config.Config) (err error) { - session := cfg.Db.NewSession() - defer session.Close() - results, _ := session.Query("select * from ip where ip in (select ip from ip group by ip having count(ip) > 1) and hostname is null order by updated desc;") +func QueryInfo(ip string) (query QueryIP, err error) { + client := http.Client{} + var url = fmt.Sprintf("https://ip.paulbsd.com/%s", ip) + req, _ := http.NewRequest("GET", url, nil) + req.Header.Add("Accept", "*/*") + req.Header.Add("User-Agent", "ipbl") + res, _ := client.Do(req) + data, _ := io.ReadAll(res.Body) + err = json.Unmarshal(data, &query) + return +} + +func Cleanup(session *xorm.Session) (err error) { + results, _ := session.Query(` + SELECT * + FROM ip + WHERE ip IN ( + SELECT ip + FROM ip + GROUP BY ip + HAVING COUNT(ip) > 1) + AND hostname IS NULL + ORDER BY updated DESC;`) if len(results) > 0 { - _, err := session.Query("delete from ip where ip in (select ip from ip group by ip having count(ip) > 1) and hostname is null;") + _, err := session.Query(` + DELETE FROM ip + WHERE ip IN ( + SELECT ip + FROM ip + GROUP by ip + HAVING COUNT(ip) > 1) + AND hostname IS NULL;`) if err != nil { log.Println("error deleting orphans") } @@ -166,33 +215,29 @@ func Cleanup(cfg *config.Config) (err error) { func (ip *IP) APIFormat() *APIIP { if ip == nil { - return nil + return &APIIP{} } return &APIIP{ - IP: ip.IP, - Rdns: ip.Rdns.String, - Src: ip.Src, - Hostname: ip.Hostname.String, - City: ip.City.String, - Country: ip.Country.String, - Date: ip.Updated.Local().String(), + IP: ip.IP, + Rdns: ip.Rdns.String, + APIAS: *ip.AutonomousSystem.APIFormat(), + APICity: ip.City.APIFormat().CityName, + APICountry: ip.Country.APIFormat().CountryName, } } -func (ip *APIIP) APIConvert() *IP { - if ip == nil { - return nil - } - return &IP{ - IP: ip.IP, +func (ip *IP) APIParse(apiip APIIP) (err error) { + *ip = IP{ + IP: apiip.IP, Rdns: sql.NullString{ - String: ip.Rdns, - Valid: false}, - Src: ip.Src, - Hostname: sql.NullString{ - String: ip.Hostname, + String: apiip.Rdns, Valid: true}, + AutonomousSystem: &AutonomousSystem{ + ASID: apiip.APIAS.ASID, + ASName: apiip.APIAS.ASName, + }, } + return } func (ip *APIIP) BeforeInsert() (err error) { @@ -200,24 +245,35 @@ func (ip *APIIP) BeforeInsert() (err error) { } type IP struct { - ID int `xorm:"pk autoincr"` - IP string `xorm:"text notnull unique(ipsrc)" json:"ip"` - Rdns sql.NullString `xorm:"text default"` - Src string `xorm:"text notnull unique(ipsrc)" json:"src"` - Hostname sql.NullString `xorm:"text default '' unique(ipsrc)" json:"hostname"` - City sql.NullString `xorm:"text default '' unique(ipsrc)" json:"city"` - Country sql.NullString `xorm:"text default '' unique(ipsrc)" json:"country"` - Created time.Time `xorm:"created notnull"` - Updated time.Time `xorm:"updated notnull"` + ID int `xorm:"pk autoincr"` + IP string `xorm:"text notnull unique"` + Rdns sql.NullString `xorm:"text default ''"` + AutonomousSystem *AutonomousSystem `xorm:"as_id int index default NULL"` + City *City `xorm:"city_id int index default NULL"` + Country *Country `xorm:"country_id int index default NULL"` + Whois string `xorm:"text index default null"` + Created time.Time `xorm:"created notnull"` + Updated time.Time `xorm:"updated notnull"` } type APIIP struct { - IP string `json:"ip"` - Rdns string `json:"rdns"` - Src string `json:"src"` - Hostname string `json:"hostname"` - City string `json:"city"` - Country string `json:"country"` - Date string `json:"date"` - Mode string `json:"mode"` + IP string `json:"ip"` + Rdns string `json:"rdns"` + APIAS APIAutonomousSystem `json:"as"` + APICity string `json:"city"` + APICountry string `json:"country"` + APIWhois string `json:"whois"` +} + +type QueryIP struct { + IP string `json:"ip"` + Rdns string `json:"hostname"` + APIAS QueryAutonomousSystem `json:"asn"` + APICity string `json:"city"` + APICountry string `json:"country"` +} + +type QueryAutonomousSystem struct { + ASID int `json:"id"` + ASName string `json:"name"` } diff --git a/src/models/models.go b/src/models/models.go new file mode 100644 index 0000000..90e1d79 --- /dev/null +++ b/src/models/models.go @@ -0,0 +1,65 @@ +package models + +import ( + "context" + "fmt" + + "git.paulbsd.com/paulbsd/ipbl/src/config" + "xorm.io/xorm" + "xorm.io/xorm/names" +) + +var ( + x *xorm.Engine + tables []interface{} + + // HasEngine specifies if we have a xorm.Engine + HasEngine bool +) + +func init() { + tables = append(tables, + new(AutonomousSystem), + new(Cfg), + new(CfgSet), + new(CfgTrustlist), + new(CfgZMQ), + new(City), + new(Country), + new(Event), + new(IP), + ) + for _, name := range []string{"SSL", "UID"} { + names.LintGonicMapper[name] = true + } +} + +// NewEngine initializes a new xorm.Engine +func NewEngine(ctx *context.Context, config *config.Config) (err error) { + var x = config.Db + + if err = x.Ping(); err != nil { + return err + } + + if err = x.Sync2(tables...); err != nil { + return fmt.Errorf("sync database struct error: %v", err) + } + + return nil +} + +// DropTables initializes a new xorm.Engine +func DropTables(ctx *context.Context, config *config.Config) (err error) { + var x = config.Db + + if err = x.Ping(); err != nil { + return err + } + + if err = x.Sync2(tables...); err != nil { + return fmt.Errorf("sync database struct error: %v", err) + } + + return nil +} diff --git a/src/models/scanresult.go b/src/models/scanresult.go new file mode 100644 index 0000000..73382fd --- /dev/null +++ b/src/models/scanresult.go @@ -0,0 +1,65 @@ +package models + +import ( + "time" + + "git.paulbsd.com/paulbsd/ipbl/src/config" +) + +func (sr *ScanResult) GetOrCreate(cfg *config.Config) (apisr *APIScanResult, err error) { + has, err := cfg.Db.Get(sr) + if !has { + cfg.Db.Insert(sr) + } else { + cfg.Db.ID(sr.ID).Update(sr) + } + sr.Get(cfg) + apisr = sr.APIFormat() + return +} + +func (sr *ScanResult) Get(cfg *config.Config) (apisr *APIScanResult, err error) { + has, err := cfg.Db.Get(sr) + if !has || err != nil { + return + } + apisr = sr.APIFormat() + return +} + +func (sr *ScanResult) APIFormat() *APIScanResult { + if sr == nil { + return &APIScanResult{} + } + return &APIScanResult{ + Protocol: sr.Protocol, + PortID: sr.PortID, + State: sr.State, + ServiceName: sr.ServiceName, + } +} + +func (sr *ScanResult) APIParse(apisr APIScanResult) (err error) { + *sr = ScanResult{ + Protocol: apisr.Protocol, + PortID: apisr.PortID} + return +} + +type ScanResult struct { + ID int `xorm:"pk autoincr"` + Protocol string `xorm:"text default ''"` + PortID int `xorm:"default ''"` + State string `xorm:"text default ''"` + ServiceName string `xorm:"text default ''"` + Created time.Time `xorm:"created notnull"` + Updated time.Time `xorm:"updated notnull"` +} + +type APIScanResult struct { + ID int `json:"id"` + Protocol string `json:"proto"` + PortID int `json:"port_id"` + State string `json:"state"` + ServiceName string `json:"service"` +} diff --git a/src/routers/funcs.go b/src/routers/funcs.go index 3361a01..4c22c44 100644 --- a/src/routers/funcs.go +++ b/src/routers/funcs.go @@ -23,22 +23,26 @@ func RegisterRoutes(e *echo.Echo, ctx *context.Context, cfg *config.Config) { }) e.GET("/ip/:ip", func(c echo.Context) (err error) { - ret, err := models.GetIP(ctx, cfg, c.Param("ip")) + session := cfg.Db.NewSession() + defer session.Close() + var ip = models.IP{IP: c.Param("ip")} + ret, err := ip.Get(session) return Result(c, err, ret) }) e.POST("/ip", func(c echo.Context) (err error) { + session := cfg.Db.NewSession() + var apiip = new(models.APIIP) var ip = new(models.IP) var msg = "No IP inserted" - err = c.Bind(ip) + err = c.Bind(apiip) if err != nil { return Result(c, fmt.Errorf("error when parsing body"), "") } - numinsert, numupdate, err := ip.InsertOrUpdate(cfg) + ip.APIParse(*apiip) + _, err = ip.GetOrCreate(session) if err != nil { return Result(c, err, "") } - msg = fmt.Sprintf("zmq: Inserted %d IP, Updated %d IP", numinsert, numupdate) - log.Println(msg) return Result(c, err, msg) }) e.GET("/ips", func(c echo.Context) (err error) { @@ -62,6 +66,8 @@ func RegisterRoutes(e *echo.Echo, ctx *context.Context, cfg *config.Config) { return Result(c, err, ret) }) e.POST("/ips", func(c echo.Context) (err error) { + session := cfg.Db.NewSession() + defer session.Close() var apiips = []models.APIIP{} var ips = []models.IP{} var msg string @@ -70,23 +76,39 @@ func RegisterRoutes(e *echo.Echo, ctx *context.Context, cfg *config.Config) { return Result(c, err, apiips) } for _, apiip := range apiips { - ips = append(ips, *apiip.APIConvert()) + var ip = new(models.IP) + ip.APIParse(apiip) + ips = append(ips, *ip) } - numinsert, numupdate, _ := models.InsertIPBulk(cfg, &ips) + numinsert, numupdate, _ := models.InsertIPBulk(session, &ips) msg = fmt.Sprintf("zmq: Inserted %d IP, Updated %d IP", numinsert, numupdate) log.Println(msg) return Result(c, err, msg) }) + e.POST("/event", func(c echo.Context) (err error) { + session := cfg.Db.NewSession() + defer session.Close() + var apievent = new(models.APIEvent) + var event = new(models.Event) + err = c.Bind(apievent) + if err != nil { + return Result(c, fmt.Errorf("error when parsing body"), "") + } + event.APIParse(session, *apievent) + err = event.Insert(cfg) + if err != nil { + return Result(c, err, "") + } + return Result(c, err, "OK") + }) e.POST("/hostinfo", func(c echo.Context) (err error) { var hostinfos = []models.HostInfo{} err = c.Bind(&hostinfos) if err != nil { return Result(c, err, hostinfos) } - num, err := cfg.Db.Insert(&hostinfos) - msg := fmt.Sprintf("Inserted %v host info (%d)", hostinfos, num) - log.Println(msg) - return Result(c, err, msg) + models.ProcessHostInfo(cfg, hostinfos) + return Result(c, err, "") }) e.GET("/config", func(c echo.Context) (err error) { trustlists, err := models.GetGlobalConfig(*cfg) diff --git a/src/routers/main.go b/src/routers/main.go index 2d5734a..6216614 100644 --- a/src/routers/main.go +++ b/src/routers/main.go @@ -24,6 +24,9 @@ func RunServer(ctx *context.Context, cfg *config.Config) (err error) { func Result(c echo.Context, inputerr error, data interface{}) (err error) { if inputerr != nil { + if inputerr.Error() == "not found" { + return c.String(http.StatusNotFound, "Content not found") + } if inputerr.Error() == "Not Found" { return c.String(http.StatusNotFound, "Content not found") } diff --git a/src/zmqrouter/main.go b/src/zmqrouter/main.go index de2c92d..5bb4915 100644 --- a/src/zmqrouter/main.go +++ b/src/zmqrouter/main.go @@ -30,44 +30,46 @@ func Init(cfg *config.Config) (err error) { func Handle(cfg *config.Config, reqsock *goczmq.Sock, pubsock *goczmq.Sock, channel string) (err error) { log.Println("Start handling zmq sockets") + var lastip string for { var msg = "err" var req, err = reqsock.RecvMessage() if err != nil { - log.Println("unable to receive message from req socket") + log.Println("unable to receive message from req socket", err) continue } - var lastip string var topub [][]byte for _, val := range req { - var apiip = models.APIIP{} - var ip = models.IP{} + var apievent = models.APIEvent{} + var event = models.Event{} - err = json.Unmarshal(val, &apiip) + err = json.Unmarshal(val, &apievent) if err != nil { log.Println("unable to parse ip address", err) continue } - if apiip.IP == lastip { - lastip = apiip.IP + if apievent.IP != "" && apievent.IP == lastip { continue } - if apiip.Mode != "init" { - ip = *apiip.APIConvert() + if apievent.Mode != "init" { + session := cfg.Db.NewSession() + event.APIParse(session, apievent) + session.Close() - numinsert, numupdate, err := ip.InsertOrUpdate(cfg) + err := event.Insert(cfg) if err != nil { log.Println(err) } - log.Printf("zmq: Inserted %d IP, Updated %d IP\n", numinsert, numupdate) + log.Printf("zmq: Inserted event") } tmpval := fmt.Sprintf("%s %s", channel, string(val)) ipjson := []byte(tmpval) topub = append(topub, ipjson) + lastip = apievent.IP } err = pubsock.SendMessage(topub)