feat(tasks): add templates findings workflows and batch controls
This commit is contained in:
1 parent
b7e0b93023
commit
d3183dbbd2
65 files changed
+11603
-2198
No files matched your search
+195
-45
@@ -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))
|
||||
|
||||
Reference in new issue
Block a user