From a81b1ac9ab4184488ed7e4fe933129f3e6f730a7 Mon Sep 17 00:00:00 2001
From: Prateek Sen <psen@newrelic.com>
Date: Fri, 26 Jul 2024 17:18:12 +0530
Subject: [PATCH] dynamodb instrmentation added

---
 lib/newrelic_security/constants.rb            |   2 +
 .../dynamodb/chain.rb                         |  52 +++++++++
 .../dynamodb/instrumentation.rb               | 104 ++++++++++++++++++
 .../dynamodb/prepend.rb                       |  39 +++++++
 4 files changed, 197 insertions(+)
 create mode 100644 lib/newrelic_security/instrumentation-security/dynamodb/chain.rb
 create mode 100644 lib/newrelic_security/instrumentation-security/dynamodb/instrumentation.rb
 create mode 100644 lib/newrelic_security/instrumentation-security/dynamodb/prepend.rb

diff --git a/lib/newrelic_security/constants.rb b/lib/newrelic_security/constants.rb
index 30c778c..49ad58b 100644
--- a/lib/newrelic_security/constants.rb
+++ b/lib/newrelic_security/constants.rb
@@ -19,6 +19,7 @@ module NewRelic::Security
   COLON_IAST_COLON = ':IAST:'
   NOSQL_DB_COMMAND = 'NOSQL_DB_COMMAND'
   SQL_DB_COMMAND = 'SQL_DB_COMMAND'
+  DYNAMO_DB_COMMAND = 'DYNAMO_DB_COMMAND'
   FILE_OPERATION = 'FILE_OPERATION'
   FILE_INTEGRITY = 'FILE_INTEGRITY'
   SYSTEM_COMMAND = 'SYSTEM_COMMAND'
@@ -30,6 +31,7 @@ module NewRelic::Security
   SQLITE = 'SQLITE'
   MYSQL = 'MYSQL'
   POSTGRES = 'POSTGRES'
+  DQL = 'DQL'
   ISO_8859_1 = 'ISO-8859-1'
   UTF_8 = 'UTF-8'
   RAILS = 'rails'
diff --git a/lib/newrelic_security/instrumentation-security/dynamodb/chain.rb b/lib/newrelic_security/instrumentation-security/dynamodb/chain.rb
new file mode 100644
index 0000000..53b4e0a
--- /dev/null
+++ b/lib/newrelic_security/instrumentation-security/dynamodb/chain.rb
@@ -0,0 +1,52 @@
+module NewRelic::Security
+  module Instrumentation
+    module Aws
+      module DynamoDB
+        module Client
+          module Chain
+
+            def self.instrument!
+              ::Aws::DynamoDB::Client.class_eval do
+                include NewRelic::Security::Instrumentation::Aws::DynamoDB::Client
+  
+                alias_method :put_item_without_security, :put_item
+  
+                def put_item(*args)
+                  retval = nil
+                  event = put_item_on_enter(*args) { retval = put_item_without_security(*args) }
+                  put_item_on_exit(event) { return retval }
+                end
+
+                alias_method :get_item_without_security, :get_item
+  
+                def get_item(*args)
+                  retval = nil
+                  event = get_item_on_enter(*args) { retval = get_item_without_security(*args) }
+                  get_item_on_exit(event) { return retval }
+                end
+
+                alias_method :update_item_without_security, :update_item
+  
+                def update_item(*args)
+                  retval = nil
+                  event = update_item_on_enter(*args) { retval = update_item_without_security(*args) }
+                  update_item_on_exit(event) { return retval }
+                end
+
+                alias_method :delete_item_without_security, :delete_item
+  
+                def delete_item(*args)
+                  retval = nil
+                  event = delete_item_on_enter(*args) { retval = delete_item_without_security(*args) }
+                  delete_item_on_exit(event) { return retval }
+                end
+                
+              end
+            end
+          end
+        end
+      end
+
+    end
+  end
+end
\ No newline at end of file
diff --git a/lib/newrelic_security/instrumentation-security/dynamodb/instrumentation.rb b/lib/newrelic_security/instrumentation-security/dynamodb/instrumentation.rb
new file mode 100644
index 0000000..e195c8d
--- /dev/null
+++ b/lib/newrelic_security/instrumentation-security/dynamodb/instrumentation.rb
@@ -0,0 +1,104 @@
+require_relative 'prepend'
+require_relative 'chain'
+
+module NewRelic::Security
+  module Instrumentation
+    module Aws::DynamoDB::Client
+      def put_item_on_enter(*args)
+        event = nil
+        NewRelic::Security::Agent.logger.debug "OnEnter : #{self.class}.#{__method__}"
+        hash = {}
+        hash[:payloadType] = :write
+        hash[:payload] = args[0]
+        event = NewRelic::Security::Agent::Control::Collector.collect(DYNAMO_DB_COMMAND, [hash], DQL)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+        return event
+      end
+
+      def put_item_on_exit(event)
+        NewRelic::Security::Agent.logger.debug "OnExit :  #{self.class}.#{__method__}"
+        NewRelic::Security::Agent::Utils.create_exit_event(event)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+      end
+
+      def get_item_on_enter(*args)
+        event = nil
+        NewRelic::Security::Agent.logger.debug "OnEnter : #{self.class}.#{__method__}"
+        hash = {}
+        hash[:payloadType] = :read
+        hash[:payload] = args[0]
+        event = NewRelic::Security::Agent::Control::Collector.collect(DYNAMO_DB_COMMAND, [hash], DQL)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+        return event
+      end
+
+      def get_item_on_exit(event)
+        NewRelic::Security::Agent.logger.debug "OnExit :  #{self.class}.#{__method__}"
+        NewRelic::Security::Agent::Utils.create_exit_event(event)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+      end
+
+      def update_item_on_enter(*args)
+        event = nil
+        NewRelic::Security::Agent.logger.debug "OnEnter : #{self.class}.#{__method__}"
+        hash = {}
+        hash[:payloadType] = :update
+        hash[:payload] = args[0]
+        event = NewRelic::Security::Agent::Control::Collector.collect(DYNAMO_DB_COMMAND, [hash], DQL)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+        return event
+      end
+
+      def update_item_on_exit(event)
+        NewRelic::Security::Agent.logger.debug "OnExit :  #{self.class}.#{__method__}"
+        NewRelic::Security::Agent::Utils.create_exit_event(event)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+      end
+
+      def delete_item_on_enter(*args)
+        event = nil
+        NewRelic::Security::Agent.logger.debug "OnEnter : #{self.class}.#{__method__}"
+        hash = {}
+        hash[:payloadType] = :delete
+        hash[:payload] = args[0]
+        event = NewRelic::Security::Agent::Control::Collector.collect(DYNAMO_DB_COMMAND, [hash], DQL)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+        return event
+      end
+
+      def delete_item_on_exit(event)
+        NewRelic::Security::Agent.logger.debug "OnExit :  #{self.class}.#{__method__}"
+        NewRelic::Security::Agent::Utils.create_exit_event(event)
+      rescue => exception
+        NewRelic::Security::Agent.logger.error "Exception in hook in #{self.class}.#{__method__}, #{exception.inspect}, #{exception.backtrace}"
+      ensure
+        yield
+      end
+
+    end
+
+  end
+end
+
+NewRelic::Security::Instrumentation::InstrumentationLoader.install_instrumentation(:dynamodb, ::Aws::DynamoDB::Client, ::NewRelic::Security::Instrumentation::Aws::DynamoDB::Client)
diff --git a/lib/newrelic_security/instrumentation-security/dynamodb/prepend.rb b/lib/newrelic_security/instrumentation-security/dynamodb/prepend.rb
new file mode 100644
index 0000000..64a6693
--- /dev/null
+++ b/lib/newrelic_security/instrumentation-security/dynamodb/prepend.rb
@@ -0,0 +1,39 @@
+module NewRelic::Security
+  module Instrumentation
+    module Aws
+      module DynamoDB
+        module Client
+          module Prepend
+            include NewRelic::Security::Instrumentation::Aws::DynamoDB::Client
+
+            def put_item(*args)
+              retval = nil
+              event = put_item_on_enter(*args) { retval = super }
+              put_item_on_exit(event) { return retval }
+            end
+
+            def get_item(*args)
+              retval = nil
+              event = get_item_on_enter(*args) { retval = super }
+              get_item_on_exit(event) { return retval }
+            end
+
+            def update_item(*args)
+              retval = nil
+              event = update_item_on_enter(*args) { retval = super }
+              update_item_on_exit(event) { return retval }
+            end
+
+            def delete_item(*args)
+              retval = nil
+              event = delete_item_on_enter(*args) { retval = super }
+              delete_item_on_exit(event) { return retval }
+            end
+
+          end
+        end
+      end
+
+    end
+  end
+end