From 62db392860353d389effc6ded27a7d53401dbd80 Mon Sep 17 00:00:00 2001 From: Issif Date: Thu, 5 Dec 2019 18:19:49 +0100 Subject: [PATCH] add Snapshots + Images ressources + fix : ressources without Region predicate was added at each run --- README.md | 2 ++ availibiltyzones.go | 11 +++++-- dgraph.go | 45 ++++++++++++++------------- images.go | 75 ++++++++++++++++++++++++++++++++------------- main.go | 21 +++++++------ schema.go | 5 +++ securitygroups.go | 13 ++++++-- snapshots.go | 12 ++++++-- 8 files changed, 124 insertions(+), 60 deletions(-) diff --git a/README.md b/README.md index 2984b5a..c3fc75e 100644 --- a/README.md +++ b/README.md @@ -174,6 +174,7 @@ Here the list of currently supported ressources : - DbInstance - DbParameterGroup - DbSubnetGroup +- Image - Instance - InstanceProfile - KeyPair @@ -183,6 +184,7 @@ Here the list of currently supported ressources : - NatGateway - OptionGroup - SecurityGroup +- Snapshot - Subnet - TargetGroup - Volume diff --git a/availibiltyzones.go b/availibiltyzones.go index 1afaa25..4440445 100644 --- a/availibiltyzones.go +++ b/availibiltyzones.go @@ -41,10 +41,18 @@ func (list availabilityZoneList) addNodes(c *connector) { return } log.Println("Add AvailibilityZones Nodes") + m := make(map[string]availabilityZoneNodes) + json.Unmarshal(c.dgraphQuery("AvailabilityZone"), &m) + a := make(availabilityZoneNodes, 0, len(list.AvailabilityZones)) for _, i := range list.AvailabilityZones { var b availabilityZoneNode + for _, j := range m["list"] { + if *i.ZoneName == j.ZoneName { + b.UID = j.UID + } + } b.Service = "ec2" b.Type = []string{"AvailabilityZone"} b.Region = c.awsRegion @@ -56,9 +64,8 @@ func (list availabilityZoneList) addNodes(c *connector) { c.dgraphAddNodes(a) c.stats.NumberOfNodes += len(a) - m := make(map[string]availabilityZoneNodes) - n := make(map[string]string) json.Unmarshal(c.dgraphQuery("AvailabilityZone"), &m) + n := make(map[string]string) for _, i := range m["list"] { n[i.ZoneName] = i.UID } diff --git a/dgraph.go b/dgraph.go index 9a0b58a..054efdd 100644 --- a/dgraph.go +++ b/dgraph.go @@ -20,35 +20,38 @@ func (c *connector) dgraphDropAll() { } func (c *connector) dgraphDropPrevious() { - log.Println("List previous data") - txn := c.dgraphClient.NewTxn() - defer txn.Discard(*c.context) + log.Println("Drop previous data") - q := `query query($owner: string, $region: string){ - list(func: eq(OwnerId, $owner)) @filter(eq(Region, $region)) { + for { + txn := c.dgraphClient.NewTxn() + q := `query query($owner: string, $region: string){ + list(func: eq(OwnerId, $owner), first: 250) @filter(eq(Region, $region)) { uid } }` - res, err := txn.QueryWithVars(*c.context, q, map[string]string{"$owner": c.awsAccountID, "$region": c.awsRegion}) - if err != nil { - log.Println(err.Error()) - } + res, err := txn.QueryWithVars(*c.context, q, map[string]string{"$owner": c.awsAccountID, "$region": c.awsRegion}) + if err != nil { + log.Println(err.Error()) + } - m := make(map[string]cidrNodes) - json.Unmarshal(res.Json, &m) + m := make(map[string]cidrNodes) // Cidr are simplest ressources, we use them for an easy Marshal + json.Unmarshal(res.Json, &m) - if len(m["list"]) != 0 { - log.Println("Drop previous data") - n, _ := json.Marshal(m["list"]) - mu := &api.Mutation{ - CommitNow: true, - DeleteJson: n, - } + if len(m["list"]) != 0 { + n, _ := json.Marshal(m["list"]) + mu := &api.Mutation{ + // CommitNow: true, + DeleteJson: n, + } - _, err = txn.Mutate(*c.context, mu) - if err != nil { - log.Fatal(err) + _, err = txn.Mutate(*c.context, mu) + if err != nil { + log.Fatal(err) + } + txn.Commit(*c.context) + } else { + break } } } diff --git a/images.go b/images.go index b7c1c35..339bd0c 100644 --- a/images.go +++ b/images.go @@ -12,25 +12,26 @@ type imageList struct{ *ec2.DescribeImagesOutput } type imageNodes []imageNode type imageNode struct { - UID string `json:"uid,omitempty"` - Type []string `json:"dgraph.type,omitempty"` - Name string `json:"name,omitempty"` // This field is only for Ratel Viz - OwnerID string `json:"OwnerId,omitempty"` - OwnerName string `json:"OwnerName,omitempty"` - Region string `json:"Region,omitempty"` - Service string `json:"Service,omitempty"` - ImageID string `json:"ImageId,omitempty"` - VirtualizationType string `json:"VirtualizationType,omitempty"` - Hypervisor string `json:"Hypervisor,omitempty"` - EnaSupport bool `json:"EnaSupport,omitempty"` - SriovNetSupport string `json:"SriovNetSupport,omitempty"` - State string `json:"State,omitempty"` - Architecture string `json:"Architecture,omitempty"` - ImageLocation string `json:"ImageLocation,omitempty"` - RootDeviceType string `json:"RootDeviceType,omitempty"` - RootDeviceName string `json:"RootDeviceName,omitempty"` - Public bool `json:"Public,omitempty"` - ImageType string `json:"ImageType,omitempty"` + UID string `json:"uid,omitempty"` + Type []string `json:"dgraph.type,omitempty"` + Name string `json:"name,omitempty"` // This field is only for Ratel Viz + OwnerID string `json:"OwnerId,omitempty"` + OwnerName string `json:"OwnerName,omitempty"` + Region string `json:"Region,omitempty"` + Service string `json:"Service,omitempty"` + ImageID string `json:"ImageId,omitempty"` + VirtualizationType string `json:"VirtualizationType,omitempty"` + Hypervisor string `json:"Hypervisor,omitempty"` + EnaSupport bool `json:"EnaSupport,omitempty"` + SriovNetSupport string `json:"SriovNetSupport,omitempty"` + State string `json:"State,omitempty"` + Architecture string `json:"Architecture,omitempty"` + ImageLocation string `json:"ImageLocation,omitempty"` + RootDeviceType string `json:"RootDeviceType,omitempty"` + RootDeviceName string `json:"RootDeviceName,omitempty"` + Public bool `json:"Public,omitempty"` + ImageType string `json:"ImageType,omitempty"` + Snapshot snapshotNodes `json:"_Snapshot,omitempty"` } func (c *connector) listImages() imageList { @@ -90,10 +91,16 @@ func (list imageList) addNodes(c *connector) { b.Public = *i.Public b.ImageType = *i.ImageType a = append(a, b) + if len(a) == 100 { + c.dgraphAddNodes(a) + c.stats.NumberOfNodes += len(a) + a = imageNodes{} + } + } + if len(a) != 0 { + c.dgraphAddNodes(a) + c.stats.NumberOfNodes += len(a) } - - c.dgraphAddNodes(a) - c.stats.NumberOfNodes += len(a) m := make(map[string]imageNodes) n := make(map[string]string) @@ -103,3 +110,27 @@ func (list imageList) addNodes(c *connector) { } c.ressources["Images"] = n } + +func (list imageList) addEdges(c *connector) { + defer c.waitGroup.Done() + + if len(list.Images) == 0 { + return + } + log.Println("Add Image Edges") + a := imageNodes{} + for _, i := range list.Images { + b := imageNode{ + UID: c.ressources["Images"][*i.ImageId], + } + for _, j := range i.BlockDeviceMappings { + if j.Ebs != nil { + if j.Ebs.SnapshotId != nil { + b.Snapshot = append(b.Snapshot, snapshotNode{UID: c.ressources["Snapshots"][*j.Ebs.SnapshotId], DeviceName: *j.DeviceName}) + } + } + } + a = append(a, b) + } + c.dgraphAddNodes(a) +} diff --git a/main.go b/main.go index 8a6ae7f..3782ec8 100644 --- a/main.go +++ b/main.go @@ -67,11 +67,11 @@ func main() { var cachesubnetgroups cacheSubnetGroupList var vpcpeeringconnections vpcPeeringConnectionList var natgateways natGatewayList - // var images imageList - // var snapshots snapshotList + var snapshots snapshotList + var images imageList // List ressources - connector.waitGroup.Add(25) + connector.waitGroup.Add(27) start := time.Now() go func() { instances = connector.listInstances() }() go func() { keypairs = connector.listKeyPairs() }() @@ -98,13 +98,13 @@ func main() { go func() { cachesubnetgroups = connector.listCacheSubnetGroups() }() go func() { vpcpeeringconnections = connector.listVpcPeeringConnections() }() go func() { natgateways = connector.listNatGateways() }() - // images = connector.listImages() - // snapshots = connector.listSnapshots() + go func() { snapshots = connector.listSnapshots() }() + go func() { images = connector.listImages() }() connector.waitGroup.Wait() // Add Nodes - connector.waitGroup.Add(25) + connector.waitGroup.Add(27) instances.addNodes(connector) keypairs.addNodes(connector) volumes.addNodes(connector) @@ -130,13 +130,13 @@ func main() { cachesubnetgroups.addNodes(connector) vpcpeeringconnections.addNodes(connector) natgateways.addNodes(connector) - // images.addNodes(connector) - // snapshots.addNodes(connector) + snapshots.addNodes(connector) + images.addNodes(connector) connector.waitGroup.Wait() // Add Edges - connector.waitGroup.Add(17) + connector.waitGroup.Add(19) instances.addEdges(connector) addresses.addEdges(connector) volumes.addEdges(connector) @@ -154,7 +154,8 @@ func main() { cachesubnetgroups.addEdges(connector) vpcpeeringconnections.addEdges(connector) natgateways.addEdges(connector) - // snapshots.addEdges(connector) + snapshots.addEdges(connector) + images.addEdges(connector) connector.waitGroup.Wait() diff --git a/schema.go b/schema.go index 1a9827b..6b6b428 100644 --- a/schema.go +++ b/schema.go @@ -23,6 +23,7 @@ func getdgraphSchema() string { _Image: uid @reverse . _Subnet: [uid] @reverse . _SecurityGroup: [uid] @reverse . + _Snapshot: [uid] @reverse . InstanceId: string @index(term) . InstanceType: string @index(term) . @@ -80,6 +81,7 @@ func getdgraphSchema() string { Device: string _AvailabilityZone: AvailabilityZone _Instance: Instance + _Snapshot: Snapshot } PublicIp: string @index(term) . @@ -296,6 +298,7 @@ func getdgraphSchema() string { Encrypted: bool . Progress: string @index(term) . _Volume: uid @reverse . + DeviceName: string @index(term) . type Snapshot { name: string @@ -309,6 +312,7 @@ func getdgraphSchema() string { VolumeSize: int Encrypted: bool Progress: string + DeviceName: string _Volume: Volume } @@ -343,6 +347,7 @@ func getdgraphSchema() string { RootDeviceName: string Public: bool ImageType: string + _Snapshot: Snapshot } SubnetId: string @index(term) . diff --git a/securitygroups.go b/securitygroups.go index 6e45520..c83eeef 100644 --- a/securitygroups.go +++ b/securitygroups.go @@ -110,13 +110,21 @@ func (list securityGroupList) addNodes(c *connector) { c.dgraphAddNodes(a) c.stats.NumberOfNodes += len(a) + o := make(map[string]cidrNodes) + json.Unmarshal(c.dgraphQuery("Cidr"), &o) if len(y) != 0 { for i := range y { - z = append(z, cidrNode{ + b := cidrNode{ Service: "ec2", Type: []string{"Cidr"}, Name: i, - }) + } + for _, j := range o["list"] { + if i == j.Name { + b.UID = j.UID + } + } + z = append(z, b) } c.dgraphAddNodes(z) c.stats.NumberOfNodes += len(z) @@ -136,7 +144,6 @@ func (list securityGroupList) addNodes(c *connector) { } c.ressources["SecurityGroupEdges"] = l - o := make(map[string]cidrNodes) p := make(map[string]string) json.Unmarshal(c.dgraphQuery("Cidr"), &o) for _, i := range o["list"] { diff --git a/snapshots.go b/snapshots.go index 05ab5f4..9d0e1d3 100644 --- a/snapshots.go +++ b/snapshots.go @@ -25,6 +25,7 @@ type snapshotNode struct { Encrypted bool `json:"Encrypted,omitempty"` Progress string `json:"Progress,omitempty"` Volume volumeNode `json:"_Volume,omitempty"` + DeviceName string `json:"_Snapshot|DeviceName,omitempty"` } func (c *connector) listSnapshots() snapshotList { @@ -73,9 +74,16 @@ func (list snapshotList) addNodes(c *connector) { b.Encrypted = *i.Encrypted b.Progress = *i.Progress a = append(a, b) + if len(a) == 100 { + c.dgraphAddNodes(a) + c.stats.NumberOfNodes += len(a) + a = snapshotNodes{} + } + } + if len(a) != 0 { + c.dgraphAddNodes(a) + c.stats.NumberOfNodes += len(a) } - c.dgraphAddNodes(a) - c.stats.NumberOfNodes += len(a) m := make(map[string]snapshotNodes) n := make(map[string]string)