phishingclub/backend/repository/proxy.go

package repository

import (
	"context"
	"fmt"
	"strings"

	"github.com/google/uuid"
	"github.com/oapi-codegen/nullable"
	"github.com/phishingclub/phishingclub/database"
	"github.com/phishingclub/phishingclub/errs"
	"github.com/phishingclub/phishingclub/model"
	"github.com/phishingclub/phishingclub/vo"
	"gorm.io/gorm"
)

var proxyAllowedColumns = assignTableToColumns(database.PROXY_TABLE, []string{
	"created_at",
	"updated_at",
	"name",
	"start_url",
})

// ProxyOption is for eager loading
type ProxyOption struct {
	Fields []string
	*vo.QueryArgs
	WithCompany bool
}

// Proxy is a proxy repository
type Proxy struct {
	DB *gorm.DB
}

// load preloads the table relations
func (m *Proxy) load(
	options *ProxyOption,
	db *gorm.DB,
) *gorm.DB {
	if options.WithCompany {
		db = db.Joins("Company")
	}
	return db
}

// Insert inserts a proxy
func (m *Proxy) Insert(
	ctx context.Context,
	proxy *model.Proxy,
) (*uuid.UUID, error) {
	id := uuid.New()
	row := proxy.ToDBMap()
	row["id"] = id
	AddTimestamps(row)

	res := m.DB.
		Model(&database.Proxy{}).
		Create(row)

	if res.Error != nil {
		return nil, res.Error
	}
	return &id, nil
}

// GetAll gets proxies
func (m *Proxy) GetAll(
	ctx context.Context,
	companyID *uuid.UUID,
	options *ProxyOption,
) (*model.Result[model.Proxy], error) {
	result := model.NewEmptyResult[model.Proxy]()
	var dbProxies []database.Proxy
	db := m.load(options, m.DB)
	db = withCompanyIncludingNullContext(db, companyID, database.PROXY_TABLE)
	db, err := useQuery(db, database.PROXY_TABLE, options.QueryArgs, proxyAllowedColumns...)
	if err != nil {
		return result, errs.Wrap(err)
	}
	if options.Fields != nil {
		fields := assignTableToColumns(database.PROXY_TABLE, options.Fields)
		db = db.Select(strings.Join(fields, ","))
	}
	dbRes := db.
		Find(&dbProxies)

	if dbRes.Error != nil {
		return result, dbRes.Error
	}

	hasNextPage, err := useHasNextPage(db, database.PROXY_TABLE, options.QueryArgs, proxyAllowedColumns...)
	if err != nil {
		return result, errs.Wrap(err)
	}
	result.HasNextPage = hasNextPage

	for _, dbProxy := range dbProxies {
		proxy, err := ToProxy(&dbProxy)
		if err != nil {
			return result, errs.Wrap(err)
		}
		result.Rows = append(result.Rows, proxy)
	}
	return result, nil
}

// GetAllSubset gets proxies with limited data
func (m *Proxy) GetAllSubset(
	ctx context.Context,
	companyID *uuid.UUID,
	options *ProxyOption,
) (*model.Result[model.ProxyOverview], error) {
	result := model.NewEmptyResult[model.ProxyOverview]()
	var dbProxies []database.Proxy
	db := withCompanyIncludingNullContext(m.DB, companyID, database.PROXY_TABLE)
	db, err := useQuery(db, database.PROXY_TABLE, options.QueryArgs, proxyAllowedColumns...)
	if err != nil {
		return result, errs.Wrap(err)
	}
	dbRes := db.
		Select("id, created_at, updated_at, name, description, start_url, company_id").
		Find(&dbProxies)

	if dbRes.Error != nil {
		return result, dbRes.Error
	}

	hasNextPage, err := useHasNextPage(db, database.PROXY_TABLE, options.QueryArgs, proxyAllowedColumns...)
	if err != nil {
		return result, errs.Wrap(err)
	}
	result.HasNextPage = hasNextPage

	for _, dbProxy := range dbProxies {
		proxyOverview := model.ProxyOverview{
			ID:          *dbProxy.ID,
			CreatedAt:   dbProxy.CreatedAt,
			UpdatedAt:   dbProxy.UpdatedAt,
			Name:        dbProxy.Name,
			Description: dbProxy.Description,
			StartURL:    dbProxy.StartURL,
			CompanyID:   dbProxy.CompanyID,
		}
		result.Rows = append(result.Rows, &proxyOverview)
	}
	return result, nil
}

// GetAllByCompanyID gets proxies by company id
func (m *Proxy) GetAllByCompanyID(
	ctx context.Context,
	companyID *uuid.UUID,
	options *ProxyOption,
) (*model.Result[model.Proxy], error) {
	result := model.NewEmptyResult[model.Proxy]()
	var dbProxies []database.Proxy
	db := m.load(options, m.DB)
	db = whereCompany(db, database.PROXY_TABLE, companyID)
	db, err := useQuery(db, database.PROXY_TABLE, options.QueryArgs, proxyAllowedColumns...)
	if err != nil {
		return result, errs.Wrap(err)
	}
	if options.Fields != nil {
		fields := assignTableToColumns(database.PROXY_TABLE, options.Fields)
		db = db.Select(strings.Join(fields, ","))
	}
	dbRes := db.
		Find(&dbProxies)

	if dbRes.Error != nil {
		return result, dbRes.Error
	}

	hasNextPage, err := useHasNextPage(db, database.PROXY_TABLE, options.QueryArgs, proxyAllowedColumns...)
	if err != nil {
		return result, errs.Wrap(err)
	}
	result.HasNextPage = hasNextPage

	for _, dbProxy := range dbProxies {
		proxy, err := ToProxy(&dbProxy)
		if err != nil {
			return result, errs.Wrap(err)
		}
		result.Rows = append(result.Rows, proxy)
	}
	return result, nil
}

// GetByID gets proxy by id
func (m *Proxy) GetByID(
	ctx context.Context,
	id *uuid.UUID,
	options *ProxyOption,
) (*model.Proxy, error) {
	dbProxy := database.Proxy{}
	db := m.load(options, m.DB)
	result := db.
		Where(TableColumnID(database.PROXY_TABLE)+" = ?", id).
		First(&dbProxy)

	if result.Error != nil {
		return nil, result.Error
	}
	return ToProxy(&dbProxy)
}

// GetByNameAndCompanyID gets proxy by name
func (m *Proxy) GetByNameAndCompanyID(
	ctx context.Context,
	name *vo.String64,
	companyID *uuid.UUID, // can be null
	options *ProxyOption,
) (*model.Proxy, error) {
	proxy := database.Proxy{}
	db := m.load(options, m.DB)
	db = withCompanyIncludingNullContext(db, companyID, database.PROXY_TABLE)
	result := db.
		Where(
			fmt.Sprintf(
				"%s = ?",
				TableColumn(database.PROXY_TABLE, "name"),
			),
			name.String(),
		).
		First(&proxy)

	if result.Error != nil {
		return nil, result.Error
	}
	return ToProxy(&proxy)
}

// UpdateByID updates a proxy by id
func (m *Proxy) UpdateByID(
	ctx context.Context,
	id *uuid.UUID,
	proxy *model.Proxy,
) error {
	row := proxy.ToDBMap()
	AddUpdatedAt(row)
	res := m.DB.
		Model(&database.Proxy{}).
		Where("id = ?", id).
		Updates(row)

	if res.Error != nil {
		return res.Error
	}
	return nil
}

// DeleteByID deletes a proxy by id
func (m *Proxy) DeleteByID(
	ctx context.Context,
	id *uuid.UUID,
) error {
	result := m.DB.Delete(&database.Proxy{}, id)

	if result.Error != nil {
		return result.Error
	}
	return nil
}

func ToProxy(row *database.Proxy) (*model.Proxy, error) {
	id := nullable.NewNullableWithValue(*row.ID)
	companyID := nullable.NewNullNullable[uuid.UUID]()
	if row.CompanyID != nil {
		companyID.Set(*row.CompanyID)
	}
	name := nullable.NewNullableWithValue(*vo.NewString64Must(row.Name))

	description, err := vo.NewOptionalString1024(row.Description)
	if err != nil {
		return nil, errs.Wrap(err)
	}
	descriptionNullable := nullable.NewNullableWithValue(*description)

	startURL := nullable.NewNullableWithValue(*vo.NewString1024Must(row.StartURL))

	proxyConfig, err := vo.NewString1MB(row.ProxyConfig)
	if err != nil {
		return nil, errs.Wrap(err)
	}
	proxyConfigNullable := nullable.NewNullableWithValue(*proxyConfig)

	return &model.Proxy{
		ID:          id,
		CreatedAt:   row.CreatedAt,
		UpdatedAt:   row.UpdatedAt,
		CompanyID:   companyID,
		Name:        name,
		Description: descriptionNullable,
		StartURL:    startURL,
		ProxyConfig: proxyConfigNullable,
	}, nil
}