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" }