From 684d81d4c209a0cc85f9f79ecf816fd75ddffc2c Mon Sep 17 00:00:00 2001
From: Aaron Lippold <alippold@mitre.org>
Date: Tue, 25 Apr 2017 19:33:10 -0400
Subject: [PATCH] psql doesn't print headers + extra output + cconfigurable
 host + docs

Signed-off-by: Aaron Lippold <lippold@gmail.com>
---
 docs/resources/postgres_session.md.erb | 27 ++++++++++++++++----------
 lib/resources/postgres_session.rb      | 21 +++++++++++---------
 2 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/docs/resources/postgres_session.md.erb b/docs/resources/postgres_session.md.erb
index a24afca86..9c7b47097 100644
--- a/docs/resources/postgres_session.md.erb
+++ b/docs/resources/postgres_session.md.erb
@@ -10,17 +10,24 @@ Use the `postgres_session` InSpec audit resource to test SQL commands run agains
 
 A `postgres_session` resource block declares the username and password to use for the session, and then the command to be run:
 
-    sql = postgres_session('username', 'password')
+    # Create a PostgreSQL session:
+    sql = postgres_session('username', 'password', 'host')
 
+    # default values:
+    #   username: 'postgres'
+    #   host: 'localhost'
+
+    # Run an SQL query with an optional database to execute
+    sql.query('sql_query', ['database_name'])`
+
+A full example is:
+
+    sql = postgres_session('username', 'password', 'host')
     describe sql.query('SELECT * FROM pg_shadow WHERE passwd IS NULL;') do
-      its('output') { should eq('') }
+      its('output') { should eq '' }
     end
 
-where
-
-* `sql = postgres_session` declares a username and password with permission to run the query
-* `sql.query('')` contains the query to be run
-* `its('output') { should eq('') }` compares the results of the query against the expected result in the test
+where `its('output') { should eq '' }` compares the results of the query against the expected result in the test
 
 ## Matchers
 
@@ -58,9 +65,9 @@ The following examples show how to use this InSpec audit resource.
 
 ### Test the PostgreSQL shadow password
 
-    sql = postgres_session('my_user', 'password')
+    sql = postgres_session('my_user', 'password', '192.168.1.2')
 
-    describe sql.query('SELECT * FROM pg_shadow WHERE passwd IS NULL;') do
+    describe sql.query('SELECT * FROM pg_shadow WHERE passwd IS NULL;', ['testdb']) do
       its('output') { should eq('') }
     end
 
@@ -70,6 +77,6 @@ The following examples show how to use this InSpec audit resource.
                   FROM pg_language
                   WHERE lanpltrusted = \'f\'
                   AND lanname!=\'internal\'
-                  AND lanname!=\'c\';') do
+                  AND lanname!=\'c\';', ['postgres']) do
       its('output') { should eq '0' }
     end
diff --git a/lib/resources/postgres_session.rb b/lib/resources/postgres_session.rb
index bd0ae4f3a..ecc44bfb7 100644
--- a/lib/resources/postgres_session.rb
+++ b/lib/resources/postgres_session.rb
@@ -2,6 +2,7 @@
 # copyright: 2015, Vulcano Security GmbH
 # author: Dominik Richter
 # author: Christoph Hartmann
+# author: Aaron Lippold
 # license: All rights reserved
 
 module Inspec::Resources
@@ -26,16 +27,23 @@ module Inspec::Resources
     name 'postgres_session'
     desc 'Use the postgres_session InSpec audit resource to test SQL commands run against a PostgreSQL database.'
     example "
-      sql = postgres_session('username', 'password')
+      sql = postgres_session('username', 'password', 'host')
+      query('sql_query', ['database_name'])` contains the query and (optional) database to execute
+
+      # default values:
+      # username: 'postgres'
+      # host: 'localhost'
+      # db: databse == db_user running the sql query
 
       describe sql.query('SELECT * FROM pg_shadow WHERE passwd IS NULL;') do
         its('output') { should eq('') }
       end
     "
 
-    def initialize(user, pass)
+    def initialize(user, pass, host = nil)
       @user = user || 'postgres'
       @pass = pass
+      @host = host || 'localhost'
     end
 
     def query(query, db = [])
@@ -44,7 +52,7 @@ module Inspec::Resources
       # that does this securely
       escaped_query = query.gsub(/\\/, '\\\\').gsub(/"/, '\\"').gsub(/\$/, '\\$')
       # run the query
-      cmd = inspec.command("PGPASSWORD='#{@pass}' psql -U #{@user} #{dbs} -h localhost -c \"#{escaped_query}\"")
+      cmd = inspec.command("PGPASSWORD='#{@pass}' psql -U #{@user} #{dbs} -h #{@host} -A -t -c \"#{escaped_query}\"")
       out = cmd.stdout + "\n" + cmd.stderr
       if cmd.exit_status != 0 or
          out =~ /could not connect to .*/ or
@@ -52,12 +60,7 @@ module Inspec::Resources
         # skip this test if the server can't run the query
         skip_resource "Can't read run query #{query.inspect} on postgres_session: #{out}"
       else
-        # remove the whole header (i.e. up to the first ^-----+------+------$)
-        # remove the tail
-        lines = cmd.stdout
-                   .sub(/(.*\n)+([-]+[+])*[-]+\n/, '')
-                   .sub(/\n[^\n]*\n\n$/, '')
-        Lines.new(lines.strip, "PostgreSQL query: #{query}")
+        Lines.new(cmd.stdout.strip, "PostgreSQL query: #{query}")
       end
     end
   end