nettopo-go/internal/scans/postgres_store.go

package scans

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"sort"
	"time"
)

type PostgresStore struct {
	db *sql.DB
}

func NewPostgresStore(db *sql.DB) *PostgresStore {
	return &PostgresStore{db: db}
}

func (s *PostgresStore) CreateScan(req CreateScanRequest) (ScanJob, error) {
	if err := validateCreateScanRequest(req); err != nil {
		return ScanJob{}, err
	}
	applyDefaultOptions(&req.Options)

	id, err := newID()
	if err != nil {
		return ScanJob{}, err
	}

	now := time.Now().UTC()
	job := ScanJob{
		ID:                id,
		Name:              req.Name,
		Status:            "queued",
		CIDRs:             req.CIDRs,
		ExcludeIPs:        req.ExcludeIPs,
		SNMPCredentialsID: req.SNMPCredentialsID,
		Options:           req.Options,
		Progress:          0,
		Stats:             ScanStats{},
		CreatedAt:         now,
	}

	cidrsJSON, err := json.Marshal(job.CIDRs)
	if err != nil {
		return ScanJob{}, err
	}
	excludeJSON, err := json.Marshal(job.ExcludeIPs)
	if err != nil {
		return ScanJob{}, err
	}
	optionsJSON, err := json.Marshal(job.Options)
	if err != nil {
		return ScanJob{}, err
	}
	statsJSON, err := json.Marshal(job.Stats)
	if err != nil {
		return ScanJob{}, err
	}

	var snmpID any
	if job.SNMPCredentialsID != "" {
		snmpID = job.SNMPCredentialsID
	}

	query := `
insert into scan_jobs (
  id, name, status, cidrs, exclude_ips, options, progress, stats, snmp_credentials_id, created_at
) values (
  $1, $2, $3, $4::jsonb, $5::jsonb, $6::jsonb, $7, $8::jsonb, $9, $10
)`

	_, err = s.db.ExecContext(context.Background(), query,
		job.ID,
		job.Name,
		job.Status,
		string(cidrsJSON),
		string(excludeJSON),
		string(optionsJSON),
		job.Progress,
		string(statsJSON),
		snmpID,
		job.CreatedAt,
	)
	if err != nil {
		return ScanJob{}, err
	}

	return job, nil
}

func (s *PostgresStore) GetScan(id string) (ScanJob, bool) {
	query := `
select id, name, status, cidrs::text, exclude_ips::text, options::text,
       progress, stats::text, coalesce(snmp_credentials_id::text, ''),
       created_at, started_at, finished_at
from scan_jobs
where id = $1
limit 1`

	var job ScanJob
	var cidrsRaw string
	var excludeRaw string
	var optionsRaw string
	var statsRaw string
	var startedAt sql.NullTime
	var finishedAt sql.NullTime

	err := s.db.QueryRowContext(context.Background(), query, id).Scan(
		&job.ID,
		&job.Name,
		&job.Status,
		&cidrsRaw,
		&excludeRaw,
		&optionsRaw,
		&job.Progress,
		&statsRaw,
		&job.SNMPCredentialsID,
		&job.CreatedAt,
		&startedAt,
		&finishedAt,
	)
	if errors.Is(err, sql.ErrNoRows) {
		return ScanJob{}, false
	}
	if err != nil {
		return ScanJob{}, false
	}

	if err = json.Unmarshal([]byte(cidrsRaw), &job.CIDRs); err != nil {
		return ScanJob{}, false
	}
	if err = json.Unmarshal([]byte(excludeRaw), &job.ExcludeIPs); err != nil {
		return ScanJob{}, false
	}
	if err = json.Unmarshal([]byte(optionsRaw), &job.Options); err != nil {
		return ScanJob{}, false
	}
	if err = json.Unmarshal([]byte(statsRaw), &job.Stats); err != nil {
		return ScanJob{}, false
	}
	if startedAt.Valid {
		job.StartedAt = startedAt.Time
	}
	if finishedAt.Valid {
		job.FinishedAt = finishedAt.Time
	}

	return job, true
}

func (s *PostgresStore) ListScans(limit int) []ScanJob {
	if limit <= 0 {
		limit = 50
	}
	query := `
select id, name, status, cidrs::text, exclude_ips::text, options::text, progress,
       stats::text, coalesce(snmp_credentials_id::text, ''), created_at, started_at, finished_at
from scan_jobs
order by created_at desc
limit $1`

	rows, err := s.db.QueryContext(context.Background(), query, limit)
	if err != nil {
		return nil
	}
	defer rows.Close()

	jobs := make([]ScanJob, 0, limit)
	for rows.Next() {
		var job ScanJob
		var cidrsRaw string
		var excludeRaw string
		var optionsRaw string
		var statsRaw string
		var startedAt sql.NullTime
		var finishedAt sql.NullTime

		err = rows.Scan(
			&job.ID,
			&job.Name,
			&job.Status,
			&cidrsRaw,
			&excludeRaw,
			&optionsRaw,
			&job.Progress,
			&statsRaw,
			&job.SNMPCredentialsID,
			&job.CreatedAt,
			&startedAt,
			&finishedAt,
		)
		if err != nil {
			return nil
		}

		if json.Unmarshal([]byte(cidrsRaw), &job.CIDRs) != nil ||
			json.Unmarshal([]byte(excludeRaw), &job.ExcludeIPs) != nil ||
			json.Unmarshal([]byte(optionsRaw), &job.Options) != nil ||
			json.Unmarshal([]byte(statsRaw), &job.Stats) != nil {
			return nil
		}
		if startedAt.Valid {
			job.StartedAt = startedAt.Time
		}
		if finishedAt.Valid {
			job.FinishedAt = finishedAt.Time
		}
		jobs = append(jobs, job)
	}

	sort.Slice(jobs, func(i, j int) bool {
		return jobs[i].CreatedAt.After(jobs[j].CreatedAt)
	})
	return jobs
}

func (s *PostgresStore) UpdateScan(id string, update ScanUpdate) (ScanJob, bool) {
	current, ok := s.GetScan(id)
	if !ok {
		return ScanJob{}, false
	}

	if update.Status != "" {
		current.Status = update.Status
	}
	if update.Progress >= 0 {
		current.Progress = update.Progress
	}
	current.Stats = update.Stats
	if update.StartedAt != nil {
		current.StartedAt = *update.StartedAt
	}
	if update.FinishedAt != nil {
		current.FinishedAt = *update.FinishedAt
	}

	optionsJSON, err := json.Marshal(current.Options)
	if err != nil {
		return ScanJob{}, false
	}
	statsJSON, err := json.Marshal(current.Stats)
	if err != nil {
		return ScanJob{}, false
	}

	var startedAt any
	if !current.StartedAt.IsZero() {
		startedAt = current.StartedAt
	}
	var finishedAt any
	if !current.FinishedAt.IsZero() {
		finishedAt = current.FinishedAt
	}

	query := `
update scan_jobs
set status = $2,
    options = $3::jsonb,
    progress = $4,
    stats = $5::jsonb,
    started_at = $6,
    finished_at = $7
where id = $1`
	_, err = s.db.ExecContext(
		context.Background(),
		query,
		id,
		current.Status,
		string(optionsJSON),
		current.Progress,
		string(statsJSON),
		startedAt,
		finishedAt,
	)
	if err != nil {
		return ScanJob{}, false
	}

	return current, true
}

func (s *PostgresStore) SaveHostResult(scanID string, host HostResult) error {
	query := `
insert into scan_hosts (scan_id, ip, is_up, checked_at)
values ($1, $2, $3, $4)`
	_, err := s.db.ExecContext(context.Background(), query, scanID, host.IP, host.IsUp, host.CheckedAt)
	return err
}

func (s *PostgresStore) ListHostResults(scanID string, limit int) []HostResult {
	if limit <= 0 {
		limit = 500
	}
	query := `
select ip::text, is_up, checked_at
from scan_hosts
where scan_id = $1
order by checked_at desc
limit $2`
	rows, err := s.db.QueryContext(context.Background(), query, scanID, limit)
	if err != nil {
		return nil
	}
	defer rows.Close()

	out := make([]HostResult, 0, limit)
	for rows.Next() {
		var h HostResult
		if err = rows.Scan(&h.IP, &h.IsUp, &h.CheckedAt); err != nil {
			return nil
		}
		out = append(out, h)
	}
	return out
}

func (s *PostgresStore) SaveOpenPortResult(scanID string, result OpenPortResult) error {
	query := `
insert into scan_ports (scan_id, ip, port, is_open, checked_at)
values ($1, $2, $3, $4, $5)`
	_, err := s.db.ExecContext(
		context.Background(),
		query,
		scanID,
		result.IP,
		result.Port,
		result.IsOpen,
		result.CheckedAt,
	)
	return err
}

func (s *PostgresStore) ListOpenPortResults(scanID string, limit int) []OpenPortResult {
	if limit <= 0 {
		limit = 1000
	}
	query := `
select ip::text, port, is_open, checked_at
from scan_ports
where scan_id = $1
order by checked_at desc
limit $2`
	rows, err := s.db.QueryContext(context.Background(), query, scanID, limit)
	if err != nil {
		return nil
	}
	defer rows.Close()

	out := make([]OpenPortResult, 0, limit)
	for rows.Next() {
		var p OpenPortResult
		if err = rows.Scan(&p.IP, &p.Port, &p.IsOpen, &p.CheckedAt); err != nil {
			return nil
		}
		out = append(out, p)
	}
	return out
}

func (s *PostgresStore) SaveSNMPResult(scanID string, result SNMPResult) error {
	query := `
insert into scan_snmp (scan_id, ip, success, sys_name, sys_descr, sys_object_id, error_text, checked_at)
values ($1, $2, $3, $4, $5, $6, $7, $8)`
	_, err := s.db.ExecContext(
		context.Background(),
		query,
		scanID,
		result.IP,
		result.Success,
		result.SysName,
		result.SysDescr,
		result.SysObjectID,
		result.Error,
		result.CheckedAt,
	)
	return err
}

func (s *PostgresStore) ListSNMPResults(scanID string, limit int) []SNMPResult {
	if limit <= 0 {
		limit = 1000
	}
	query := `
select ip::text, success, coalesce(sys_name, ''), coalesce(sys_descr, ''), coalesce(sys_object_id, ''),
       coalesce(error_text, ''), checked_at
from scan_snmp
where scan_id = $1
order by checked_at desc
limit $2`
	rows, err := s.db.QueryContext(context.Background(), query, scanID, limit)
	if err != nil {
		return nil
	}
	defer rows.Close()

	out := make([]SNMPResult, 0, limit)
	for rows.Next() {
		var sres SNMPResult
		if err = rows.Scan(
			&sres.IP,
			&sres.Success,
			&sres.SysName,
			&sres.SysDescr,
			&sres.SysObjectID,
			&sres.Error,
			&sres.CheckedAt,
		); err != nil {
			return nil
		}
		out = append(out, sres)
	}
	return out
}

func (s *PostgresStore) SaveLLDPResult(scanID string, result LLDPResult) error {
	query := `
insert into scan_lldp (scan_id, ip, local_port_num, remote_chassis_id, remote_port_id, remote_sys_name, checked_at)
values ($1, $2, $3, $4, $5, $6, $7)`
	_, err := s.db.ExecContext(
		context.Background(),
		query,
		scanID,
		result.IP,
		result.LocalPortNum,
		result.RemoteChassisID,
		result.RemotePortID,
		result.RemoteSysName,
		result.CheckedAt,
	)
	return err
}

func (s *PostgresStore) ListLLDPResults(scanID string, limit int) []LLDPResult {
	if limit <= 0 {
		limit = 2000
	}
	query := `
select ip::text, coalesce(local_port_num, ''), coalesce(remote_chassis_id, ''),
       coalesce(remote_port_id, ''), coalesce(remote_sys_name, ''), checked_at
from scan_lldp
where scan_id = $1
order by checked_at desc
limit $2`
	rows, err := s.db.QueryContext(context.Background(), query, scanID, limit)
	if err != nil {
		return nil
	}
	defer rows.Close()

	out := make([]LLDPResult, 0, limit)
	for rows.Next() {
		var lr LLDPResult
		if err = rows.Scan(
			&lr.IP,
			&lr.LocalPortNum,
			&lr.RemoteChassisID,
			&lr.RemotePortID,
			&lr.RemoteSysName,
			&lr.CheckedAt,
		); err != nil {
			return nil
		}
		out = append(out, lr)
	}
	return out
}