feat(tasks): add templates findings workflows and batch controls

This commit is contained in:
pphua committed 2026-08-21 21:01:39 +08:00
1 parent b7e0b93023
commit d3183dbbd2
65 files changed
+11603 -2198

No files matched your search

+195 -45
View File
@@ -83,6 +83,7 @@ type PortService struct {
type AssetStore struct {
db *DB
company *CompanyStore
tx *sql.Tx
}
// Assets returns the asset store.
@@ -93,6 +94,39 @@ func (d *DB) Assets() *AssetStore {
// Companies returns the company store associated with this asset store.
func (s *AssetStore) Companies() *CompanyStore { return s.company }
// withCompanyScopeMutation serializes scope resolution and every asset write
// that consumes its result in one transaction. Nested asset side effects reuse
// the same transaction through the scoped store.
func (s *AssetStore) withCompanyScopeMutation(fn func(*AssetStore) (int64, error)) (int64, error) {
if s.tx != nil {
return fn(s)
}
tx, err := s.db.Begin()
if err != nil {
return 0, err
}
defer tx.Rollback() //nolint:errcheck
if err := lockCompanyScopeMutation(tx); err != nil {
return 0, err
}
scoped := &AssetStore{db: s.db, company: s.company, tx: tx}
id, err := fn(scoped)
if err != nil {
return 0, err
}
if err := tx.Commit(); err != nil {
return 0, err
}
return id, nil
}
func (s *AssetStore) resolveCompanyWithICP(rootDomain, ipStr, icp string) (*int64, error) {
if s.tx != nil {
return resolveCompanyWithICP(s.tx, rootDomain, ipStr, icp)
}
return s.company.ResolveCompanyWithICP(rootDomain, ipStr, icp)
}
// =====================================================================
// Helpers
// =====================================================================
@@ -219,7 +253,12 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
companyID, err := s.company.ResolveCompany(domain, "")
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertRootDomain(req)
})
}
companyID, err := s.resolveCompanyWithICP(domain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -237,12 +276,20 @@ func (s *AssetStore) UpsertRootDomain(req UpsertRootDomainReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, $4::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, icp, company_id, company_source, task_ids)
VALUES ('root_domain', $1, $1, $2, $3, 'scope', $4::bigint[])
ON CONFLICT (domain) WHERE type = 'root_domain' DO UPDATE SET
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
extra = assets.extra || EXCLUDED.extra,
last_seen = now()
@@ -267,8 +314,13 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
if req.IP == "" {
return 0, fmt.Errorf("ip is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertIP(req)
})
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany("", req.IP)
companyID, err := s.resolveCompanyWithICP("", req.IP, "")
if err != nil {
return 0, err
}
@@ -292,9 +344,9 @@ func (s *AssetStore) UpsertIP(req UpsertIPReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, ip, c_segment, bound_domains, open_ports, company_id, company_source, task_ids)
VALUES ('ip', $1, $2::cidr, $3::text[], $4::jsonb[], $5, 'scope', $6::bigint[])
ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
bound_domains = (SELECT ARRAY(SELECT DISTINCT unnest(assets.bound_domains || EXCLUDED.bound_domains))),
open_ports = (
@@ -306,7 +358,15 @@ ON CONFLICT (ip) WHERE type = 'ip' DO UPDATE SET
)
),
c_segment = COALESCE(assets.c_segment, EXCLUDED.c_segment),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.IP, csegVal, boundDomains, openPortsJSON, companyID, taskIDs).Scan(&id)
@@ -367,14 +427,19 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
if domain == "" {
return 0, fmt.Errorf("domain is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertSubdomain(req)
})
}
rootDomain, _ := RootDomain(domain)
if rootDomain == "" {
rootDomain = domain
}
// resolve company by root domain
companyID, err := s.company.ResolveCompany(rootDomain, "")
// Resolve by root domain first, then the asset's exact normalized ICP.
companyID, err := s.resolveCompanyWithICP(rootDomain, "", req.ICP)
if err != nil {
return 0, err
}
@@ -409,7 +474,7 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
cseg := calcCSegment(ipStr)
if companyID == nil && ipStr != "" {
companyID, _ = s.company.ResolveCompany("", ipStr)
companyID, _ = s.resolveCompanyWithICP("", ipStr, "")
}
var taskIDs string
@@ -432,14 +497,22 @@ func (s *AssetStore) UpsertSubdomain(req UpsertSubdomainReq) (id int64, err erro
recordType := req.RecordType
recordValueArr := marshalStringArray(req.RecordValue)
err = s.db.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, $9::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, domain, root_domain, record_type, record_value, ip, c_segment, icp, company_id, company_source, task_ids)
VALUES ('subdomain', $1, $2, $3, $4::text[], $5, $6::cidr, $7, $8, 'scope', $9::bigint[])
ON CONFLICT (domain, COALESCE(record_type,'')) WHERE type = 'subdomain' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
icp = COALESCE(EXCLUDED.icp, assets.icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
record_value = (SELECT ARRAY(SELECT DISTINCT unnest(assets.record_value || EXCLUDED.record_value))),
extra = assets.extra || EXCLUDED.extra,
@@ -459,7 +532,7 @@ type UpsertAppReq struct {
Category string
Description string
ICP string
CompanyID *int64 // explicit override; nil = no auto attribution for apps
CompanyID *int64 // explicit override; nil = exact ICP auto-attribution when available
TaskID int64
}
@@ -468,6 +541,11 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
if req.Name == "" {
return 0, fmt.Errorf("app name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertApp(req)
})
}
var taskIDs string
if req.TaskID > 0 {
@@ -477,6 +555,7 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
var bundleVal, catVal, descVal, icpVal, companyIDVal any
companySource := "scope"
if req.BundleID != "" {
bundleVal = req.BundleID
}
@@ -491,36 +570,65 @@ func (s *AssetStore) UpsertApp(req UpsertAppReq) (int64, error) {
}
if req.CompanyID != nil {
companyIDVal = *req.CompanyID
companySource = "explicit"
} else {
companyID, err := s.resolveCompanyWithICP("", "", req.ICP)
if err != nil {
return 0, err
}
if companyID != nil {
companyIDVal = *companyID
}
}
var id int64
var err error
if req.BundleID != "" {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', $1, $2, $3, $4, $5, $6, $7, $8::bigint[])
ON CONFLICT (bundle_id) WHERE type = 'app' AND bundle_id IS NOT NULL DO UPDATE SET
app_name = COALESCE(EXCLUDED.app_name, assets.app_name),
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, bundleVal, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
} else {
err = s.db.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6::bigint[])
err = s.tx.QueryRow(`
INSERT INTO assets(type, bundle_id, app_name, category, app_description, app_icp, company_id, company_source, task_ids)
VALUES ('app', NULL, $1, $2, $3, $4, $5, $6, $7::bigint[])
ON CONFLICT (app_name) WHERE type = 'app' AND bundle_id IS NULL DO UPDATE SET
category = COALESCE(EXCLUDED.category, assets.category),
app_description = COALESCE(EXCLUDED.app_description, assets.app_description),
app_icp = COALESCE(EXCLUDED.app_icp, assets.app_icp),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN EXCLUDED.company_id
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN EXCLUDED.company_source = 'explicit' THEN 'explicit'
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
last_seen = now()
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, taskIDs).Scan(&id)
RETURNING id`, req.Name, catVal, descVal, icpVal, companyIDVal, companySource, taskIDs).Scan(&id)
}
return id, err
}
@@ -548,6 +656,11 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
if req.URL == "" {
return 0, fmt.Errorf("url is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertHTTPService(req)
})
}
normURL := normalizeURL(req.URL)
domain, port, serviceName := parseURL(normURL)
rootDomain, _ := RootDomain(domain)
@@ -556,7 +669,7 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -592,16 +705,16 @@ func (s *AssetStore) UpsertHTTPService(req UpsertHTTPServiceReq) (int64, error)
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, service_type, service_name, domain, ip, port, root_domain,
c_segment, favicon_mmh3, technologies, status_code, content_length, page_title,
auth, company_id, task_ids
auth, company_id, company_source, task_ids
)
VALUES (
'service', $1, 'http', $2, $3, $4, $5, $6,
$7::cidr, $8, $9::text[], $10, $11, $12,
$13::jsonb[], $14, $15::bigint[]
$13::jsonb[], $14, 'scope', $15::bigint[]
)
ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
status_code = COALESCE(EXCLUDED.status_code, assets.status_code),
@@ -618,7 +731,15 @@ ON CONFLICT (url) WHERE type = 'service' AND service_type = 'http' DO UPDATE SET
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, serviceName, domainVal, ipVal, nullableInt(port), rootDomain,
@@ -675,6 +796,11 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
if req.ServiceName == "" {
return 0, fmt.Errorf("service_name is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertOtherService(req)
})
}
// normalize service_name (lowercase) so the (domain,ip,port,service_name)
// dedup key doesn't split "SSH" and "ssh" into separate rows.
serviceName := strings.ToLower(strings.TrimSpace(req.ServiceName))
@@ -689,7 +815,7 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -721,12 +847,12 @@ func (s *AssetStore) UpsertOtherService(req UpsertOtherServiceReq) (int64, error
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, service_type, service_name, domain, ip, port, root_domain,
c_segment, auth, company_id, task_ids
c_segment, auth, company_id, company_source, task_ids
)
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, $9::bigint[])
VALUES ('service', 'other', $1, $2, $3, $4, $5, $6::cidr, $7::jsonb[], $8, 'scope', $9::bigint[])
ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE type = 'service' AND service_type = 'other' DO UPDATE SET
domain = COALESCE(EXCLUDED.domain, assets.domain),
ip = COALESCE(EXCLUDED.ip, assets.ip),
@@ -738,7 +864,15 @@ ON CONFLICT (COALESCE(domain,''), COALESCE(ip,''), port, service_name) WHERE typ
)
),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
serviceName, domainVal, ipVal, req.Port, rootDomainVal,
@@ -799,6 +933,11 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
if req.Method == "" {
return 0, fmt.Errorf("method is required")
}
if s.tx == nil {
return s.withCompanyScopeMutation(func(scoped *AssetStore) (int64, error) {
return scoped.UpsertEndpoint(req)
})
}
method := strings.ToUpper(req.Method)
normURL := normalizeURL(req.URL)
domain, port, _ := parseURL(normURL)
@@ -808,7 +947,7 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
cseg := calcCSegment(req.IP)
companyID, err := s.company.ResolveCompany(rootDomain, req.IP)
companyID, err := s.resolveCompanyWithICP(rootDomain, req.IP, "")
if err != nil {
return 0, err
}
@@ -840,12 +979,12 @@ func (s *AssetStore) UpsertEndpoint(req UpsertEndpointReq) (int64, error) {
}
var id int64
err = s.db.QueryRow(`
err = s.tx.QueryRow(`
INSERT INTO assets(
type, url, method, domain, ip, port, root_domain,
params, company_id, task_ids, c_segment
params, company_id, company_source, task_ids, c_segment
)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, $9::bigint[], $10::cidr)
VALUES ('endpoint', $1, $2, $3, $4, $5, $6, $7::jsonb[], $8, 'scope', $9::bigint[], $10::cidr)
ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
params = (
SELECT ARRAY(
@@ -857,7 +996,15 @@ ON CONFLICT (url, method) WHERE type = 'endpoint' DO UPDATE SET
ip = COALESCE(EXCLUDED.ip, assets.ip),
c_segment = COALESCE(EXCLUDED.c_segment, assets.c_segment),
task_ids = (SELECT ARRAY(SELECT DISTINCT unnest(assets.task_ids || EXCLUDED.task_ids))),
company_id = COALESCE(EXCLUDED.company_id, assets.company_id),
company_id = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN assets.company_id
ELSE COALESCE(EXCLUDED.company_id, assets.company_id)
END,
company_source = CASE
WHEN assets.company_source = 'explicit' AND assets.company_id IS NOT NULL THEN 'explicit'
WHEN EXCLUDED.company_id IS NOT NULL THEN 'scope'
ELSE assets.company_source
END,
last_seen = now()
RETURNING id`,
normURL, method, domainVal, ipVal, nullableInt(port), rootDomainVal,
@@ -884,6 +1031,9 @@ func (s *AssetStore) QueryByType(typ string, limit, offset int) ([]*Asset, error
if limit <= 0 {
limit = 50
}
if offset < 0 {
offset = 0
}
rows, err := s.db.Query(`
SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(domain,''), COALESCE(root_domain,''), COALESCE(ip,''),
@@ -897,7 +1047,7 @@ SELECT id, type, company_id, array_to_json(task_ids)::text,
COALESCE(method,''), array_to_json(params)::text, extra, last_seen::text
FROM assets
WHERE type = $1
ORDER BY last_seen DESC
ORDER BY last_seen DESC, id DESC
LIMIT $2 OFFSET $3`, typ, limit, offset)
if err != nil {
return nil, err
@@ -954,7 +1104,7 @@ func (s *AssetStore) CountByCompany(companyID int64, typ string) (int, error) {
}
func pageClause(args *[]any, limit, offset int) string {
q := ` ORDER BY last_seen DESC`
q := ` ORDER BY last_seen DESC, id DESC`
if limit > 0 {
*args = append(*args, limit)
q += fmt.Sprintf(` LIMIT $%d`, len(*args))