EdgeAPI/internal/db/models/ip_library_file_dao.go

package models

import (
	"encoding/json"
	"errors"
	"github.com/TeaOSLab/EdgeAPI/internal/db/models/regions"
	"github.com/TeaOSLab/EdgeAPI/internal/utils"
	"github.com/TeaOSLab/EdgeCommon/pkg/iplibrary"
	_ "github.com/go-sql-driver/mysql"
	"github.com/iwind/TeaGo/Tea"
	"github.com/iwind/TeaGo/dbs"
	"github.com/iwind/TeaGo/types"
	"os"
	"time"
)

const (
	IPLibraryFileStateEnabled  = 1 // 已启用
	IPLibraryFileStateDisabled = 0 // 已禁用
)

type IPLibraryFileDAO dbs.DAO

func NewIPLibraryFileDAO() *IPLibraryFileDAO {
	return dbs.NewDAO(&IPLibraryFileDAO{
		DAOObject: dbs.DAOObject{
			DB:     Tea.Env,
			Table:  "edgeIPLibraryFiles",
			Model:  new(IPLibraryFile),
			PkName: "id",
		},
	}).(*IPLibraryFileDAO)
}

var SharedIPLibraryFileDAO *IPLibraryFileDAO

func init() {
	dbs.OnReady(func() {
		SharedIPLibraryFileDAO = NewIPLibraryFileDAO()
	})
}

// EnableIPLibraryFile 启用条目
func (this *IPLibraryFileDAO) EnableIPLibraryFile(tx *dbs.Tx, id uint64) error {
	_, err := this.Query(tx).
		Pk(id).
		Set("state", IPLibraryFileStateEnabled).
		Update()
	return err
}

// DisableIPLibraryFile 禁用条目
func (this *IPLibraryFileDAO) DisableIPLibraryFile(tx *dbs.Tx, id int64) error {
	_, err := this.Query(tx).
		Pk(id).
		Set("state", IPLibraryFileStateDisabled).
		Update()
	return err
}

// FindEnabledIPLibraryFile 查找启用中的条目
func (this *IPLibraryFileDAO) FindEnabledIPLibraryFile(tx *dbs.Tx, id int64) (*IPLibraryFile, error) {
	result, err := this.Query(tx).
		Pk(id).
		State(IPLibraryFileStateEnabled).
		Find()
	if result == nil {
		return nil, err
	}
	return result.(*IPLibraryFile), err
}

// CreateLibraryFile 创建文件
func (this *IPLibraryFileDAO) CreateLibraryFile(tx *dbs.Tx, template string, emptyValues []string, fileId int64, countries []string, provinces [][2]string, cities [][3]string, towns [][4]string, providers []string) (int64, error) {
	var op = NewIPLibraryFileOperator()
	op.Template = template

	if emptyValues == nil {
		emptyValues = []string{}
	}
	emptyValuesJSON, err := json.Marshal(emptyValues)
	if err != nil {
		return 0, err
	}
	op.EmptyValues = emptyValuesJSON

	op.FileId = fileId

	if countries == nil {
		countries = []string{}
	}
	countriesJSON, err := json.Marshal(countries)
	if err != nil {
		return 0, err
	}
	op.Countries = countriesJSON

	if provinces == nil {
		provinces = [][2]string{}
	}
	provincesJSON, err := json.Marshal(provinces)
	if err != nil {
		return 0, err
	}
	op.Provinces = provincesJSON

	if cities == nil {
		cities = [][3]string{}
	}
	citiesJSON, err := json.Marshal(cities)
	if err != nil {
		return 0, err
	}
	op.Cities = citiesJSON

	if towns == nil {
		towns = [][4]string{}
	}
	townsJSON, err := json.Marshal(towns)
	if err != nil {
		return 0, err
	}
	op.Towns = townsJSON

	if providers == nil {
		providers = []string{}
	}
	providersJSON, err := json.Marshal(providers)
	if err != nil {
		return 0, err
	}
	op.Providers = providersJSON

	op.IsFinished = false
	op.State = IPLibraryFileStateEnabled
	return this.SaveInt64(tx, op)
}

// FindAllUnfinishedLibraryFiles 查找所有未完成的文件
func (this *IPLibraryFileDAO) FindAllUnfinishedLibraryFiles(tx *dbs.Tx) (result []*IPLibraryFile, err error) {
	_, err = this.Query(tx).
		State(IPLibraryFileStateEnabled).
		Result("id", "fileId", "createdAt"). // 这里不需要其他信息
		Attr("isFinished", false).
		DescPk().
		Slice(&result).
		FindAll()
	return
}

// UpdateLibraryFileIsFinished 设置文件为已完成
func (this *IPLibraryFileDAO) UpdateLibraryFileIsFinished(tx *dbs.Tx, fileId int64) error {
	return this.Query(tx).
		Pk(fileId).
		Set("isFinished", true).
		UpdateQuickly()
}

// FindLibraryFileCountries 获取IP库中的国家/地区
func (this *IPLibraryFileDAO) FindLibraryFileCountries(tx *dbs.Tx, fileId int64) ([]string, error) {
	countriesJSON, err := this.Query(tx).
		Result("countries").
		Pk(fileId).
		FindJSONCol()
	if err != nil {
		return nil, err
	}

	if IsNull(countriesJSON) {
		return nil, nil
	}

	var result = []string{}
	err = json.Unmarshal(countriesJSON, &result)
	if err != nil {
		return nil, err
	}

	return result, nil
}

// FindLibraryFileProvinces 获取IP库中的省份
func (this *IPLibraryFileDAO) FindLibraryFileProvinces(tx *dbs.Tx, fileId int64) ([][2]string, error) {
	provincesJSON, err := this.Query(tx).
		Result("provinces").
		Pk(fileId).
		FindJSONCol()
	if err != nil {
		return nil, err
	}

	if IsNull(provincesJSON) {
		return nil, nil
	}

	var result = [][2]string{}
	err = json.Unmarshal(provincesJSON, &result)
	if err != nil {
		return nil, err
	}

	return result, nil
}

// FindLibraryFileCities 获取IP库中的城市
func (this *IPLibraryFileDAO) FindLibraryFileCities(tx *dbs.Tx, fileId int64) ([][3]string, error) {
	citiesJSON, err := this.Query(tx).
		Result("cities").
		Pk(fileId).
		FindJSONCol()
	if err != nil {
		return nil, err
	}

	if IsNull(citiesJSON) {
		return nil, nil
	}

	var result = [][3]string{}
	err = json.Unmarshal(citiesJSON, &result)
	if err != nil {
		return nil, err
	}

	return result, nil
}

// FindLibraryFileProviders 获取IP库中的ISP运营商
func (this *IPLibraryFileDAO) FindLibraryFileProviders(tx *dbs.Tx, fileId int64) ([]string, error) {
	providersJSON, err := this.Query(tx).
		Result("providers").
		Pk(fileId).
		FindJSONCol()
	if err != nil {
		return nil, err
	}

	if IsNull(providersJSON) {
		return nil, nil
	}

	var result = []string{}
	err = json.Unmarshal(providersJSON, &result)
	if err != nil {
		return nil, err
	}

	return result, nil
}

func (this *IPLibraryFileDAO) GenerateIPLibrary(tx *dbs.Tx, libraryFileId int64) error {
	one, err := this.Query(tx).Pk(libraryFileId).Find()
	if err != nil {
		return err
	}
	if one == nil {
		return errors.New("the library file not found")
	}

	var libraryFile = one.(*IPLibraryFile)
	template, err := iplibrary.NewTemplate(libraryFile.Template)
	if err != nil {
		return errors.New("create template from '" + libraryFile.Template + "' failed: " + err.Error())
	}

	var fileId = int64(libraryFile.FileId)
	if fileId == 0 {
		return errors.New("the library file has not been uploaded yet")
	}

	var dir = Tea.Root + "/data"
	stat, err := os.Stat(dir)

	if err != nil {
		if os.IsNotExist(err) {
			err = os.Mkdir(dir, 0777)
			if err != nil {
				return errors.New("can not open dir '" + dir + "' to write: " + err.Error())
			}
		} else {
			return errors.New("can not open dir '" + dir + "' to write: " + err.Error())
		}
	} else if !stat.IsDir() {
		_ = os.Remove(dir)

		err = os.Mkdir(dir, 0777)
		if err != nil {
			return errors.New("can not open dir '" + dir + "' to write: " + err.Error())
		}
	}

	// TODO 删除以往生成的文件，但要考虑到文件正在被别的任务所使用

	// 国家
	var countries = []*iplibrary.Country{}
	// TODO

	// 省份
	var provinces = []*iplibrary.Province{}
	// TODO

	// 城市
	var cities = []*iplibrary.City{}
	// TODO

	// 区县
	var towns = []*iplibrary.Town{}
	// TODO

	// ISP运营商
	var providers = []*iplibrary.Provider{}
	// TODO

	var libraryCode = utils.Sha1RandomString() // 每次都生成新的code
	var filePath = dir + "/" + this.composeFilename(libraryFileId, libraryCode)
	writer, err := iplibrary.NewFileWriter(filePath, &iplibrary.Meta{
		Author:    "", // 将来用户可以自行填写
		CreatedAt: time.Now().Unix(),
		Countries: countries,
		Provinces: provinces,
		Cities:    cities,
		Towns:     towns,
		Providers: providers,
	})
	if err != nil {
		return err
	}

	defer func() {
		_ = writer.Close()
	}()

	err = writer.WriteMeta()
	if err != nil {
		return errors.New("write meta failed: " + err.Error())
	}

	chunkIds, err := SharedFileChunkDAO.FindAllFileChunkIds(tx, fileId)
	if err != nil {
		return err
	}

	// countries etc ...
	var countryMap = map[string]int64{} // countryName => countryId
	dbCountries, err := regions.SharedRegionCountryDAO.FindAllCountries(tx)
	if err != nil {
		return err
	}
	for _, country := range dbCountries {
		for _, code := range country.AllCodes() {
			countryMap[code] = int64(country.Id)
		}
	}

	var provinceMap = map[string]int64{} // countryId_provinceName => provinceId
	dbProvinces, err := regions.SharedRegionProvinceDAO.FindAllEnabledProvinces(tx)
	if err != nil {
		return err
	}
	for _, province := range dbProvinces {
		for _, code := range province.AllCodes() {
			provinceMap[types.String(province.CountryId)+"_"+code] = int64(province.Id)
		}
	}

	var cityMap = map[string]int64{} // provinceId_cityName => cityId
	dbCities, err := regions.SharedRegionCityDAO.FindAllEnabledCities(tx)
	if err != nil {
		return err
	}
	for _, city := range dbCities {
		for _, code := range city.AllCodes() {
			cityMap[types.String(city.ProvinceId)+"_"+code] = int64(city.Id)
		}
	}

	var townMap = map[string]int64{} // cityId_townName => townId
	// TODO 将来实现区县

	var providerMap = map[string]int64{} // providerName => providerId
	dbProviders, err := regions.SharedRegionProviderDAO.FindAllEnabledProviders(tx)
	if err != nil {
		return err
	}
	for _, provider := range dbProviders {
		for _, code := range provider.AllCodes() {
			providerMap[code] = int64(provider.Id)
		}
	}

	dataParser, err := iplibrary.NewParser(&iplibrary.ParserConfig{
		Template:    template,
		EmptyValues: libraryFile.DecodeEmptyValues(),
		Iterator: func(values map[string]string) error {
			var ipFrom = values["ipFrom"]
			var ipTo = values["ipTo"]

			var countryName = values["country"]
			var provinceName = values["province"]
			var cityName = values["city"]
			var townName = values["town"]
			var providerName = values["provider"]

			var countryId = countryMap[countryName]
			var provinceId int64
			var cityId int64 = 0
			var townId int64 = 0
			var providerId = providerMap[providerName]

			if countryId > 0 {
				provinceId = provinceMap[types.String(countryId)+"_"+provinceName]
				if provinceId > 0 {
					cityId = cityMap[types.String(provinceId)+"_"+cityName]
					if cityId > 0 {
						townId = townMap[types.String(cityId)+"_"+townName]
					}
				}
			}

			err = writer.Write(ipFrom, ipTo, countryId, provinceId, cityId, townId, providerId)
			if err != nil {
				return errors.New("write failed: " + err.Error())
			}

			return nil
		},
	})
	if err != nil {
		return err
	}

	for _, chunkId := range chunkIds {
		chunk, err := SharedFileChunkDAO.FindFileChunk(tx, chunkId)
		if err != nil {
			return err
		}
		if chunk == nil {
			return errors.New("invalid chunk file, please upload again")
		}
		dataParser.Write(chunk.Data)
		err = dataParser.Parse()
		if err != nil {
			return err
		}
	}

	err = writer.Close()
	if err != nil {
		return err
	}

	// 设置code
	err = this.Query(tx).
		Pk(libraryFileId).
		Set("code", libraryCode).
		Set("isFinished", true).
		Set("generatedAt", time.Now().Unix()).
		UpdateQuickly()
	if err != nil {
		return err
	}

	return nil
}

// 组合IP库文件名
func (this *IPLibraryFileDAO) composeFilename(libraryFileId int64, code string) string {
	return "ip-library-" + types.String(libraryFileId) + "-" + code + ".db"
}