// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn . package iplibrary import ( "bytes" "encoding/json" "fmt" "github.com/TeaOSLab/EdgeAdmin/internal/utils/sizes" "github.com/TeaOSLab/EdgeAdmin/internal/web/actions/actionutils" "github.com/TeaOSLab/EdgeAdmin/internal/web/actions/default/settings/ip-library/iplibraryutils" iplib "github.com/TeaOSLab/EdgeCommon/pkg/iplibrary" "github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb" "github.com/iwind/TeaGo/Tea" "github.com/iwind/TeaGo/actions" "github.com/iwind/TeaGo/logs" "io" "os" "path/filepath" "strings" ) type UploadAction struct { actionutils.ParentAction } func (this *UploadAction) Init() { this.Nav("", "", "upload") } func (this *UploadAction) RunGet(params struct{}) { this.Data["canAccess"] = iplibraryutils.CanAccess() this.Show() } func (this *UploadAction) RunPost(params struct { Name string File *actions.File Must *actions.Must CSRF *actionutils.CSRF }) { if len(params.Name) == 0 { this.FailField("name", "请输入IP库名称") return } if params.File == nil { this.Fail("请选择要上传的IP库文件") return } fp, err := params.File.OriginFile.Open() if err != nil { this.Fail("读取IP库文件失败:" + err.Error()) return } defer func() { _ = fp.Close() }() data, err := io.ReadAll(io.LimitReader(fp, 64*sizes.M)) // 最大不超过64M if err != nil { this.Fail("读取IP库文件失败:" + err.Error()) return } // 只支持 MaxMind 格式文件(.mmdb) filename := strings.ToLower(params.File.Filename) if !strings.HasSuffix(filename, ".mmdb") { this.Fail("只支持 MaxMind 格式文件(.mmdb),请上传 GeoLite2-City.mmdb 或 GeoLite2-ASN.mmdb 文件") return } // MaxMind 格式文件,保存到 data/iplibrary/ 目录 iplibDir := Tea.Root + "/data/iplibrary" err = os.MkdirAll(iplibDir, 0755) if err != nil { this.Fail("创建IP库目录失败:" + err.Error()) return } // 根据文件名判断是 City 还是 ASN var targetPath string if strings.Contains(filename, "city") { targetPath = filepath.Join(iplibDir, "maxmind-city.mmdb") } else if strings.Contains(filename, "asn") { targetPath = filepath.Join(iplibDir, "maxmind-asn.mmdb") } else { this.Fail("MaxMind 文件名必须包含 'city' 或 'asn'") return } // 保存文件(使用临时文件原子替换) tmpPath := targetPath + ".tmp" err = os.WriteFile(tmpPath, data, 0644) if err != nil { this.Fail("保存IP库文件失败:" + err.Error()) return } // 原子替换 err = os.Rename(tmpPath, targetPath) if err != nil { os.Remove(tmpPath) this.Fail("替换IP库文件失败:" + err.Error()) return } // 通过 RPC 将文件上传到 EdgeAPI _, err = this.RPC().IPLibraryRPC().UploadMaxMindFile(this.AdminContext(), &pb.UploadMaxMindFileRequest{ Filename: params.File.Filename, Data: data, }) if err != nil { logs.Println("[IP_LIBRARY]upload MaxMind file to EdgeAPI failed: " + err.Error()) // 继续执行,不影响本地保存 } // 创建简单的 Meta meta := &iplib.Meta{ Version: 3, Code: "maxmind", Author: "MaxMind", } meta.Init() // TODO 检查是否要自动创建省市区 // 上传IP库文件到数据库 fileResp, err := this.RPC().FileRPC().CreateFile(this.AdminContext(), &pb.CreateFileRequest{ Filename: params.File.Filename, Size: int64(len(data)), IsPublic: false, MimeType: "", Type: "ipLibraryArtifact", }) if err != nil { this.ErrorPage(err) return } var fileId = fileResp.FileId var buf = make([]byte, 256*1024) var dataReader = bytes.NewReader(data) for { n, err := dataReader.Read(buf) if n > 0 { _, chunkErr := this.RPC().FileChunkRPC().CreateFileChunk(this.AdminContext(), &pb.CreateFileChunkRequest{ FileId: fileId, Data: buf[:n], }) if chunkErr != nil { this.Fail("上传文件到数据库失败:" + chunkErr.Error()) return } } if err != nil { break } } // 创建IP库信息 metaJSON, err := json.Marshal(meta) if err != nil { this.Fail("元数据编码失败:" + err.Error()) return } createResp, err := this.RPC().IPLibraryArtifactRPC().CreateIPLibraryArtifact(this.AdminContext(), &pb.CreateIPLibraryArtifactRequest{ FileId: fileId, MetaJSON: metaJSON, Name: params.Name, }) if err != nil { this.Fail("创建IP库失败:" + err.Error()) return } // 获取所有IP库记录,将其他记录的 isPublic 设为 false,新上传的设为 true artifactsResp, err := this.RPC().IPLibraryArtifactRPC().FindAllIPLibraryArtifacts(this.AdminContext(), &pb.FindAllIPLibraryArtifactsRequest{}) if err != nil { // 如果获取列表失败,不影响上传成功,只记录日志 logs.Println("[IP_LIBRARY]failed to get all artifacts: " + err.Error()) } else { // 将所有其他记录设为未使用 for _, artifact := range artifactsResp.IpLibraryArtifacts { if artifact.Id != createResp.IpLibraryArtifactId { if artifact.IsPublic { // 将其他使用中的记录设为未使用 _, err = this.RPC().IPLibraryArtifactRPC().UpdateIPLibraryArtifactIsPublic(this.AdminContext(), &pb.UpdateIPLibraryArtifactIsPublicRequest{ IpLibraryArtifactId: artifact.Id, IsPublic: false, }) if err != nil { logs.Println("[IP_LIBRARY]failed to update artifact " + fmt.Sprintf("%d", artifact.Id) + " isPublic: " + err.Error()) } } } } // 将新上传的记录设为使用中 _, err = this.RPC().IPLibraryArtifactRPC().UpdateIPLibraryArtifactIsPublic(this.AdminContext(), &pb.UpdateIPLibraryArtifactIsPublicRequest{ IpLibraryArtifactId: createResp.IpLibraryArtifactId, IsPublic: true, }) if err != nil { logs.Println("[IP_LIBRARY]failed to set new artifact as public: " + err.Error()) } } this.Success() }