ipbl/src/models/ip.go

package models

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"runtime"
	"sync"
	"time"

	"git.paulbsd.com/paulbsd/ipbl/src/config"
	"xorm.io/xorm"
)

const SCANLIMIT = 100

const IPINFO_WS = "https://ip.paulbsd.com"

var lastday = time.Now().Add(-(time.Hour * 24))

func GetIPs(ctx *context.Context, config *config.Config, limit int) (apiips []*APIIP, err error) {
	var ips []IP
	err = config.Db.Limit(limit).Desc("created").Find(&ips)
	for _, ml := range ips {
		apiips = append(apiips, ml.APIFormat())
	}
	return
}

func GetIPsLast(ctx *context.Context, config *config.Config, interval string) (apiips []string, err error) {
	var ips []IP
	err = config.Db.Where("updated >= (now()-?::interval)", interval).GroupBy("ip").Find(&ips)
	for _, ml := range ips {
		apiips = append(apiips, ml.IP)
	}
	return
}

func (ip *IP) GetOrCreate(session *xorm.Session) (apiip *APIIP, err error) {
	ip.City.GetOrCreate(session)
	ip.Country.GetOrCreate(session)
	ip.AutonomousSystem.GetOrCreate(session)
	session.Commit()

	var tmpip *IP
	if ip.ID != 0 {
		tmpip = &IP{ID: ip.ID, IP: ip.IP}
	} else {
		tmpip = &IP{IP: ip.IP}
	}
	has, err := session.Get(tmpip)
	if err != nil {
		log.Println(err)
	}
	if !has {
		session.Insert(ip)
	} else {
		ip.ID = tmpip.ID
		session.ID(ip.ID).AllCols().Update(ip)
		session.ID(ip.ID).Cols("city_id", "country_id", "as_id").Update(ip)
	}
	session.Commit()

	ip.Get(session)
	apiip = ip.APIFormat()
	return
}

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
	}
	apiip = ip.APIFormat()
	return
}

func (ip *IP) InsertOrUpdate(session *xorm.Session) (numinsert int64, numupdate int64, err error) {
	has, err := session.Get(ip)
	if has {
		session.ID(ip.ID).Update(&IP{})
	}
	session.Commit()
	return
}

func InsertIPBulk(session *xorm.Session, ips *[]IP) (numinsert int64, numupdate int64, err error) {
	for _, ip := range *ips {
		numinsert, numupdate, err = ip.InsertOrUpdate(session)
	}
	Cleanup(session)
	return
}

func ScanIP(cfg *config.Config) (err error) {
	var numthreads int = runtime.NumCPU() / 2
	for {
		session := cfg.Db.NewSession()
		orphans := []IP{}
		err = session.Where(`
			(
				(
					as_id IS NULL
					OR country_id IS NULL
					OR city_id IS NULL
					OR rdns = ''
					OR rdns IS NULL
				)
				AND updated < now()-'1d'::interval
			)
			OR
		  (
				as_id IS NULL
				AND country_id IS NULL
				AND city_id IS NULL
				AND rdns IS NULL
			)
		`).Desc("updated").Limit(SCANLIMIT).Find(&orphans)
		session.Close()

		if err == nil && len(orphans) > 0 {
			orphanchan := make(chan IP)
			var wg sync.WaitGroup

			for thr := range make([]int, numthreads) {
				go ScanOrphan(&wg, orphanchan, thr, cfg)
			}

			for _, orphan := range orphans {
				orphanchan <- orphan
			}

			close(orphanchan)
			wg.Wait()
		} else {
			time.Sleep(2 * time.Second)
		}
	}
}

func ScanOrphan(wg *sync.WaitGroup, orphans chan IP, thr int, cfg *config.Config) (err error) {
	wg.Add(1)

	session := cfg.Db.NewSession()
	defer session.Close()
	queryclient := http.Client{}

	for {
		orphan, more := <-orphans
		if more {
			if err != nil {
				log.Println(err)
			}
			var query QueryIP
			query, err := QueryInfo(&queryclient, orphan.IP)
			if err != nil {
				log.Println(err)
				time.Sleep(1 * time.Minute)
				continue
			}

			var as = AutonomousSystem{ASID: query.APIAS.Number, ASName: query.APIAS.Org}
			orphan.AutonomousSystem = &as

			if query.APICity != "" {
				var city = City{CityName: query.APICity}
				orphan.City = &city
			}

			if query.APICountry != "" {
				var country = Country{CountryName: query.APICountry}
				orphan.Country = &country
			}

			orphan.Rdns = sql.NullString{String: query.Rdns, Valid: true}
			if cfg.Options.Debug {
				log.Printf("t%d: %s -> \"%s\"\n", thr, orphan.IP, query.Rdns)
			}

			_, err = orphan.GetOrCreate(session)
			if err != nil {
				continue
			}

		} else {
			log.Printf("All orphan migrated on thread num %d\n", thr)
			wg.Done()
			break
		}
	}
	err = session.Commit()
	if err != nil {
		log.Println(err)
	}
	return nil
}

func QueryInfo(client *http.Client, ip string) (query QueryIP, err error) {
	var url = fmt.Sprintf("%s/%s", IPINFO_WS, ip)
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Println(err)
	}
	req.Header.Add("Accept", "*/*")
	req.Header.Add("User-Agent", "ipbl")
	res, err := client.Do(req)
	if err != nil {
		log.Println(err)
		return
	}
	data, err := io.ReadAll(res.Body)
	if err != nil {
		log.Println(err)
	}
	err = json.Unmarshal(data, &query)
	if err != nil {
		log.Println(err)
	}
	return
}

func Cleanup(session *xorm.Session) (err error) {
	return
}

func (ip *IP) APIFormat() *APIIP {
	if ip == nil {
		return &APIIP{}
	}
	return &APIIP{
		IP:         ip.IP,
		Rdns:       ip.Rdns.String,
		APIAS:      *ip.AutonomousSystem.APIFormat(),
		APICity:    ip.City.APIFormat().CityName,
		APICountry: ip.Country.APIFormat().CountryName,
	}
}

func (ip *IP) APIParse(apiip APIIP) (err error) {
	*ip = IP{
		IP: apiip.IP,
		Rdns: sql.NullString{
			String: apiip.Rdns,
			Valid:  true},
		AutonomousSystem: &AutonomousSystem{
			ASID:   apiip.APIAS.ASID,
			ASName: apiip.APIAS.ASName,
		},
	}
	return
}

func (ip *APIIP) BeforeInsert() (err error) {
	return
}

type IP struct {
	ID               int               `xorm:"pk autoincr"`
	IP               string            `xorm:"text notnull unique"`
	Rdns             sql.NullString    `xorm:"text index 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"`
	Created          time.Time         `xorm:"created notnull"`
	Updated          time.Time         `xorm:"updated index notnull"`
}

type APIIP struct {
	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:"as"`
	APICity    string                `json:"city"`
	APICountry string                `json:"country"`
}

type QueryAutonomousSystem struct {
	Number int    `json:"number"`
	Org    string `json:"org"`
}