package helpers import ( "fmt" "os" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "github.com/glebarez/sqlite" "gorm.io/gorm" ) const ( CNMergeTargetUserID = 11 CNSyncedCopyOwnerUserID = 12 CNTakenUsernameUserID = 13 CNSoftAffOwnerUserID = 14 OVCreateUserID = 10000001 OVMergeUserID = 10000002 OVHardSyncedUserID = 10000003 OVUsernameConflictUserID = 10000004 OVSoftAffUserID = 10000005 OVSyncedCopyUserID = 10000100 UsernameCNMergeTarget = "cn-merge-target" UsernameCNSyncedCopyOwner = "cn-synced-copy-owner" UsernameCNTakenUsername = "taken-username" UsernameCNSoftAffOwner = "aff-owner" UsernameOVCreate = "ov-create" UsernameOVMerge = "ov-merge" UsernameOVHardSynced = "ov-hard-synced" UsernameOVUsernameConflict = "taken-username" UsernameOVSoftAff = "ov-aff-soft" UsernameOVSyncedCopy = "ov-existing-synced-copy" EmailCNMergeTarget = "merge@example.com" EmailCNSyncedCopyOwner = "synced@example.com" EmailCNTakenUsername = "taken@example.com" EmailCNSoftAffOwner = "soft-owner@example.com" EmailOVCreate = "create@example.com" EmailOVMerge = "merge@example.com" EmailOVHardSynced = "synced@example.com" EmailOVUsernameConflict = "username-conflict@example.com" EmailOVSoftAff = "soft@example.com" EmailOVSyncedCopy = "copy@example.com" AffCodeCNMergeTarget = "CNM1" AffCodeCNSyncedCopyOwner = "CNH3" AffCodeCNTakenUsername = "CNT1" AffCodeCNSoftAffOwner = "SOFT1" AffCodeOVCreate = "OVC1" AffCodeOVMerge = "OVM2" AffCodeOVHardSynced = "OVH3" AffCodeOVUsername = "OVU4" AffCodeOVSyncedCopy = "OVS1" ) func ResetDB(path string) error { if err := os.Remove(path); err != nil && !os.IsNotExist(err) { return err } return nil } func OpenSQLite(path string) (*gorm.DB, error) { db, err := gorm.Open(sqlite.Open(path), &gorm.Config{}) if err != nil { return nil, err } sqlDB, err := db.DB() if err != nil { return nil, err } sqlDB.SetMaxOpenConns(1) return db, nil } func CloseSQLite(db *gorm.DB) error { if db == nil { return nil } sqlDB, err := db.DB() if err != nil { return err } return sqlDB.Close() } func SeedCNData(dbPath string) error { db, err := OpenSQLite(dbPath) if err != nil { return fmt.Errorf("open CN sqlite: %w", err) } defer func() { _ = CloseSQLite(db) }() passwordHash, err := common.Password2Hash("fixture-password") if err != nil { return fmt.Errorf("hash fixture password: %w", err) } now := common.GetTimestamp() users := []*model.User{ { Id: CNMergeTargetUserID, Username: UsernameCNMergeTarget, Password: passwordHash, DisplayName: "CN Merge Target", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailCNMergeTarget, Group: "default", Quota: 500, AffCode: AffCodeCNMergeTarget, Source: common.UserSourceLocal, CreatedAt: now, }, { Id: CNSyncedCopyOwnerUserID, Username: UsernameCNSyncedCopyOwner, Password: passwordHash, DisplayName: "CN Synced Copy Owner", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailCNSyncedCopyOwner, Group: "default", Quota: 300, AffCode: AffCodeCNSyncedCopyOwner, Source: common.UserSourceLocal, CreatedAt: now, }, { Id: CNTakenUsernameUserID, Username: UsernameCNTakenUsername, Password: passwordHash, DisplayName: "CN Username Conflict", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailCNTakenUsername, Group: "default", Quota: 50, AffCode: AffCodeCNTakenUsername, Source: common.UserSourceLocal, CreatedAt: now, }, { Id: CNSoftAffOwnerUserID, Username: UsernameCNSoftAffOwner, Password: passwordHash, DisplayName: "CN Soft Aff Owner", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailCNSoftAffOwner, Group: "default", Quota: 80, AffCode: AffCodeCNSoftAffOwner, Source: common.UserSourceLocal, CreatedAt: now, }, } for _, user := range users { if err := db.Create(user).Error; err != nil { return fmt.Errorf("seed CN user %s: %w", user.Username, err) } } return nil } func SeedOVData(dbPath string) error { db, err := OpenSQLite(dbPath) if err != nil { return fmt.Errorf("open OV sqlite: %w", err) } defer func() { _ = CloseSQLite(db) }() passwordHash, err := common.Password2Hash("fixture-password") if err != nil { return fmt.Errorf("hash fixture password: %w", err) } now := common.GetTimestamp() users := []*model.User{ { Id: OVCreateUserID, Username: UsernameOVCreate, Password: passwordHash, DisplayName: "OV Create", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailOVCreate, Group: "default", Quota: 100, AffCode: AffCodeOVCreate, Source: common.UserSourceLocal, CreatedAt: now, }, { Id: OVMergeUserID, Username: UsernameOVMerge, Password: passwordHash, DisplayName: "OV Merge", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailOVMerge, Group: "default", Quota: 200, AffCode: AffCodeOVMerge, Source: common.UserSourceLocal, CreatedAt: now, }, { Id: OVHardSyncedUserID, Username: UsernameOVHardSynced, Password: passwordHash, DisplayName: "OV Hard Synced", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailOVHardSynced, Group: "default", Quota: 300, AffCode: AffCodeOVHardSynced, Source: common.UserSourceLocal, CreatedAt: now, }, { Id: OVUsernameConflictUserID, Username: UsernameOVUsernameConflict, Password: passwordHash, DisplayName: "OV Username Conflict", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailOVUsernameConflict, Group: "default", Quota: 120, AffCode: AffCodeOVUsername, Source: common.UserSourceLocal, CreatedAt: now, }, { Id: OVSoftAffUserID, Username: UsernameOVSoftAff, Password: passwordHash, DisplayName: "OV Soft Aff", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailOVSoftAff, Group: "default", Quota: 150, AffCode: AffCodeCNSoftAffOwner, Source: common.UserSourceLocal, CreatedAt: now, }, } for _, user := range users { if err := db.Create(user).Error; err != nil { return fmt.Errorf("seed OV user %s: %w", user.Username, err) } } return SeedOVSyncedCopy(dbPath) } func SeedOVSyncedCopy(dbPath string) error { db, err := OpenSQLite(dbPath) if err != nil { return fmt.Errorf("open OV sqlite for synced copy: %w", err) } defer func() { _ = CloseSQLite(db) }() passwordHash, err := common.Password2Hash("fixture-password") if err != nil { return fmt.Errorf("hash fixture password for synced copy: %w", err) } now := common.GetTimestamp() user := &model.User{ Id: OVSyncedCopyUserID, Username: UsernameOVSyncedCopy, Password: passwordHash, DisplayName: "OV Existing Synced Copy", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Email: EmailOVSyncedCopy, Group: "default", Quota: 0, AffCode: AffCodeOVSyncedCopy, Source: common.UserSourceSynced, RemoteUserId: CNSyncedCopyOwnerUserID, SyncedQuota: 300, LastSyncAt: now, CreatedAt: now, } if err := db.Create(user).Error; err != nil { return fmt.Errorf("seed OV synced copy %s: %w", user.Username, err) } return nil }