app/periscope/periscope_controller.rb in periscope_rails-0.0.4 vs app/periscope/periscope_controller.rb in periscope_rails-0.0.5

@@ -1,76 +1,83 @@
-class PeriscopeController < ApplicationController
-  before_filter :authenticate
-  protect_from_forgery :except => [:look, :login]
-
-  def look
-    if !params[:sql].nil?
-      render :json => run_sql(params[:sql])
-    else
-      render :json => {:error => "Command not understood"}
-    end
-  end
-
-  def login
-    render :json => get_info()
-  end
-
-  private
-
-  def authenticate
-    unless PeriscopeRails::Config.check_password(params[:password].to_s)
-      render :json => {:error => "Password invalid."}
-    end
-  end
-
-  def run_sql(sql_command)
-    #TODO: protect based on CFG, not blacklist
-    bad_words = %W{drop delete update into insert index add remove grant revoke create createdb}
-    bad_words += %W{createuser createrole destroy disconnect exec execute dropdb primary key rollback ; --}
-
-    rows = nil
-    error_message = nil
-    command = sql_command.to_s.strip
-    command_words = command.downcase.gsub(/[^a-zA-Z0-9]/, " ").gsub(/\s+/, " ").split(" ")
-    if command == ""
-      #nothing
-    elsif (command_words & bad_words).size > 0
-      error_message = "Potentially harmful keyword found, blocking script."
-    else
-      begin
-        #custom_db_creds = PeriscopeRails::Config.get_db_creds()
-        #if custom_db_creds.nil?
-        #  active_record = ActiveRecord::Base
-        #else
-        #  # TODO: Establish this connection once rather than every time a query is issued
-        #  active_record = Class.new(ActiveRecord::Base)
-        #  custom_db_config = ActiveRecord::Base.connection_config.merge(custom_db_creds)
-        #  active_record.establish_connection(custom_db_config)
-        #end
-        active_record = PeriscopeRails::Config.get_active_record()
-        active_record.transaction do
-          rows = active_record.connection.select_all(command)
-          rows.each do |row|
-            row.each_key do |column|
-              if PeriscopeRails::Config.matches_filter(column)
-                row[column] = '[FILTERED]'
-              end
-            end
-          end
-          raise "OK" #abort all transactions for extra protection
-        end
-      rescue Exception => e
-        error_message = e.message unless e.message == "OK"
-      end
-    end
-    return {:error => error_message, :data => rows}
-  end
-
-  def get_info
-    tables = []
-    table_names = ActiveRecord::Base.connection.tables.sort
-    table_names.each do |table_name|
-      tables << {:name => table_name, :columns => ActiveRecord::Base.connection.columns(table_name)}
-    end
-    return {:tables => tables, :error => nil}
-  end
-end
+class PeriscopeController < ActionController::Base
+  before_filter :authenticate
+  protect_from_forgery :except => [:look, :login]
+
+  def look
+    if !params[:sql].nil?
+      render :json => run_sql(params[:sql])
+    else
+      render :json => {:error => "Command not understood"}
+    end
+  end
+
+  def login
+    render :json => get_info()
+  end
+
+  private
+
+  def authenticate
+    unless PeriscopeRails::Config.check_password(params[:password].to_s)
+      render :json => {:error => "Password invalid."}
+    end
+  end
+
+  def run_sql(sql_command)
+    #TODO: protect based on CFG, not blacklist
+    bad_words = %W{drop delete update into insert index add remove grant revoke create createdb}
+    bad_words += %W{createuser createrole destroy disconnect exec execute dropdb primary key rollback ; --}
+
+    rows = nil
+    error_message = nil
+    command = sql_command.to_s.strip
+    command_words = command.downcase.gsub(/[^a-zA-Z0-9]/, " ").gsub(/\s+/, " ").split(" ")
+    if command == ""
+      #nothing
+    elsif (command_words & bad_words).size > 0
+      error_message = "Potentially harmful keyword found, blocking script."
+    else
+      begin
+        active_record = PeriscopeRails::Config.get_active_record()
+        active_record.transaction do
+          if PeriscopeRails::Config.block_expensive_queries?
+            cost_row = ""
+            begin
+              cost_row = active_record.connection.select_all("explain #{command}")[0]["QUERY PLAN"]
+            rescue
+              puts "Warning: Periscope was unable to cost this query (1): #{command}"
+            end
+            cost_row =~ /rows=(\d+) width=(\d+)\)$/
+            row_count, width = $1.to_i, $2.to_i
+            if row_count > 0 and width > 0
+              raise "Command blocked, it may be too slow. Estimated at #{row_count} rows, commands must return fewer than #{PeriscopeRails::Config.max_rows} rows." if row_count > PeriscopeRails::Config.max_rows
+              raise "Command blocked, it may be too slow. Estimated at #{row_count * width} bytes, commands use less than #{PeriscopeRails::Config.max_size} bytes." if row_count * width > PeriscopeRails::Config.max_size
+            else
+              puts "Warning: Periscope was unable to cost this query (2): #{command}"
+            end
+          end
+          rows = active_record.connection.select_all(command)
+          rows.each do |row|
+            row.each_key do |column|
+              if PeriscopeRails::Config.matches_filter(column)
+                row[column] = '[FILTERED]'
+              end
+            end
+          end
+          raise "OK" #abort all transactions for extra protection
+        end
+      rescue Exception => e
+        error_message = e.message unless e.message == "OK"
+      end
+    end
+    return {:error => error_message, :data => rows}
+  end
+
+  def get_info
+    tables = []
+    table_names = ActiveRecord::Base.connection.tables.sort
+    table_names.each do |table_name|
+      tables << {:name => table_name, :columns => ActiveRecord::Base.connection.columns(table_name)}
+    end
+    return {:tables => tables, :error => nil}
+  end
+end