func()

in tools/gcpviz/main.go [1038:1351]


func (v *GcpViz) EnrichAssets() error {
	// Iterate BoltDB
	tx, err := v.AssetDatabase.Begin(false)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	/* Step 1: alias generation, configured via relations file */
	aliasAssetTypes := make([]string, 0, len(v.Relations.Aliases))
	for ak, _ := range v.Relations.Aliases {
		aliasAssetTypes = append(aliasAssetTypes, ak)
	}

	fmt.Fprintf(os.Stderr, "\nCreating reference aliases for assets...\n")
	wrtx, err := v.AssetDatabase.Begin(true)
	if err != nil {
		return err
	}
	err = tx.Bucket([]byte("Assets")).ForEach(func(bk, bv []byte) error {
		isAliasAsset := false
		for _, aliasAssetType := range aliasAssetTypes {
			// Simple optimization to avoid unmarshaling tons of JSON
			if strings.Contains(string(bv), aliasAssetType) {
				isAliasAsset = true
				break
			}
		}
		if !isAliasAsset {
			return nil
		}

		asset, resource, err := v.parseJsonAsset(bv)
		if err != nil {
			return err
		}

		name := asset.GetName()
		assetType := asset.GetAssetType()
		_resource := resource.(map[string]interface{})
		if aliases, ok := v.Relations.Aliases[assetType]; ok {
			for _, jsonPath := range aliases {
				if res, ok := _resource["resource"]; ok {
					_data := res.(map[string]interface{})
					if _, ok := _data["data"]; ok {
						targets, err := jsonPath(_data)
						if err != nil {
							continue
						}

						_targets, err := v.jsonPathResultsToString(targets)
						for _, vertex := range _targets {
							err = wrtx.Bucket([]byte("Aliases")).Put([]byte(vertex), []byte(name))
							v.TotalAliases++
						}
					}
				}

			}
		}
		return nil
	})
	if err != nil {
		wrtx.Rollback()
		return err
	} else {
		if err = wrtx.Commit(); err != nil {
			return err
		}
	}

	fmt.Fprintf(os.Stderr, "Creating references between assets...\n")
	tx, err = v.AssetDatabase.Begin(false)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	relationsAssetTypes := make([]string, 0, len(v.Relations.AssetTypes))
	for rk, _ := range v.Relations.AssetTypes {
		relationsAssetTypes = append(relationsAssetTypes, rk)
	}

	ipAssetTypes := make([]string, 0, len(v.Relations.IpAddresses))
	for ik, _ := range v.Relations.IpAddresses {
		ipAssetTypes = append(ipAssetTypes, ik)
	}

	ipAddresses := make([]IpAddressLink, 0)

	/* Step 2: standard reference generation, configured via relations file */
	err = tx.Bucket([]byte("Assets")).ForEach(func(bk, bv []byte) error {
		isRelationsAsset := false
		isIpAsset := false
		for _, relationsAssetType := range relationsAssetTypes {
			// Simple optimization to avoid unmarshaling tons of JSON
			if strings.Contains(string(bv), relationsAssetType) {
				isRelationsAsset = true
				break
			}
		}
		for _, ipAssetType := range ipAssetTypes {
			// Simple optimization to avoid unmarshaling tons of JSON
			if strings.Contains(string(bv), ipAssetType) {
				isIpAsset = true
				break
			}
		}

		if !isRelationsAsset && !isIpAsset {
			return nil
		}

		asset, resource, err := v.parseJsonAsset(bv)
		if err != nil {
			return err
		}

		assetType := asset.GetAssetType()
		name := asset.GetName()

		if isIpAsset {
			_resource := resource.(map[string]interface{})
			if relations, ok := v.Relations.IpAddresses[assetType]; ok {
				for _, jsonPath := range relations {
					if res, ok := _resource["resource"]; ok {
						_data := res.(map[string]interface{})
						if _, ok := _data["data"]; ok {
							targets, err := jsonPath(_data)
							if err != nil {
								continue
							}
							_targets, err := v.jsonPathResultsToString(targets)
							for _, ipRange := range _targets {
								if ipRange == "0.0.0.0/0" {
									continue
								}
								if !strings.Contains(ipRange, "/") {
									ipRange = ipRange + "/32"
								}
								_, parsedIp, err := net.ParseCIDR(ipRange)
								if err != nil {
									fmt.Fprintf(os.Stderr, "Warning: Failed to parse IP %v for resource %v.\n", ipRange, name)
								}
								ip := IpAddressLink{Ip: parsedIp, Resource: name, AssetType: assetType}
								ipAddresses = append(ipAddresses, ip)
								v.TotalIps++
							}
						}
					}
				}
			}
		}

		if isRelationsAsset {
			_resource := resource.(map[string]interface{})
			if relations, ok := v.Relations.AssetTypes[assetType]; ok {
				for _, jsonPath := range relations {
					if res, ok := _resource["resource"]; ok {
						_data := res.(map[string]interface{})
						if _, ok := _data["data"]; ok {
							targets, err := jsonPath(_data)
							if err != nil {
								continue
							}
							_targets, err := v.jsonPathResultsToString(targets)
							for _, vertex := range _targets {
								if strings.HasPrefix(vertex, "https://www.googleapis.com/compute/v1/") {
									vertex = strings.Replace(vertex, "https://www.googleapis.com/compute/v1/", "//compute.googleapis.com/", 1)
								}
								if !strings.HasPrefix(vertex, "//") && strings.Contains(vertex, ".googleapis.com/") {
									vertex = fmt.Sprintf("//%s", vertex)
								}
								rotx, err := v.AssetDatabase.Begin(false)
								if err != nil {
									return err
								}
								defer rotx.Rollback()

								target := rotx.Bucket([]byte("Assets")).Get([]byte(vertex))
								if target == nil {
									alias := rotx.Bucket([]byte("Aliases")).Get([]byte(vertex))
									if alias != nil {
										vertex = string(alias)
									} else {
										if !strings.HasPrefix(vertex, "//") && !strings.Contains(vertex, ".googleapis.com/") {
											p := strings.SplitN(assetType, "/", 2)
											vertex = fmt.Sprintf("//%s/%s", p[0], vertex)
										}
									}
								}

								v.QW.AddQuad(cayley.Quad(name, "uses", vertex, assetType))
								v.TotalEdges++
							}
						}
					}
				}
			}
		}

		return nil
	})

	/* Step 3: Link via IP addresses */
	fmt.Fprintf(os.Stderr, "Processing resource IP addresses...\n")
	for ik, ip := range ipAddresses {
		for sik, sip := range ipAddresses {
			if ik != sik {
				if ip.AssetType != sip.AssetType && sip.Ip.Contains(ip.Ip.IP) {
					v.QW.AddQuad(cayley.Quad(ip.Resource, "uses", sip.Resource, ip.AssetType))
					v.TotalEdges++
				}
			}
		}
	}
	/* Step 4: Enrich existing asset types by incorporating linked assets into new fields */
	enrichAssetTypes := make([]string, 0, len(v.Relations.Enrich))
	for ek, _ := range v.Relations.Enrich {
		enrichAssetTypes = append(enrichAssetTypes, ek)
	}
	enrichedAssets := make(map[string]map[string][]interface{}, 0)

	fmt.Fprintf(os.Stderr, "Integrating subassets as part of main assets...\n")
	err = tx.Bucket([]byte("Assets")).ForEach(func(bk, bv []byte) error {
		isEnrichAsset := false
		for _, enrichAssetType := range enrichAssetTypes {
			// Simple optimization to avoid unmarshaling tons of JSON
			if strings.Contains(string(bv), enrichAssetType) {
				isEnrichAsset = true
				break
			}
		}
		if !isEnrichAsset {
			return nil
		}
		asset, _, err := v.parseJsonAsset(bv)
		if err != nil {
			return err
		}

		assetType := asset.GetAssetType()
		name := asset.GetName()
		qval, qerr := cquad.AsValue(name)
		if !qerr {
			return err
		}

		var subAssets []interface{}
		subAssets = append(subAssets, nil)
		for _, subAssetTypes := range v.Relations.Enrich[assetType] {
			for subAssetType, _ := range subAssetTypes {
				subAssets = append(subAssets, subAssetType)
			}
		}
		newFields := make(map[string][]interface{}, 0)

		p := cayley.StartPath(v.QS, qval).LabelContext(subAssets...).In("uses")
		p.Iterate(nil).EachValue(v.QS, func(val cquad.Value) {
			target := cquad.NativeOf(val).(string)
			targetAsset, err := v.getAsset(target)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Warning: could not fetch sub-asset %v\n", target)
			} else {
				_data := targetAsset.Resource.Data.(map[string]interface{})
				for field, subAssetTypes := range v.Relations.Enrich[assetType] {
					for subAssetType, jsonPath := range subAssetTypes {
						if subAssetType == targetAsset.AssetType {
							targets, err := jsonPath(_data)
							if err != nil {
								fmt.Fprintf(os.Stderr, "Warning: JSON-Path error in enrichment: %v\n", err)
								continue
							}
							if _, ok := newFields[field]; !ok {
								newFields[field] = make([]interface{}, 0)
							}
							newFields[field] = append(newFields[field], targets)
						}
					}
				}
				enrichedAssets[name] = newFields
			}
		})

		return nil
	})

	tx, err = v.AssetDatabase.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	for name, newFields := range enrichedAssets {
		asset, err := v.getAsset(name)
		if err == nil {
			for k, v := range newFields {
				asset.Resource.Data.(map[string]interface{})[k] = v
			}
			err = v.UpdateAsset(tx, name, asset)
			if err != nil {
				return err
			}
		} else {
			return err
		}
	}
	if err = tx.Commit(); err != nil {
		return err
	}

	fmt.Fprintf(os.Stderr, "\nTotal vertexes: %d, total edges: %d, total aliases: %d, total IPs: %d\n", v.TotalVertexes, v.TotalEdges, v.TotalAliases, v.TotalIps)
	return nil
}