From 75a76a52f4ee9bdff34e9ee2fa5ccfc9dc68cfc9 Mon Sep 17 00:00:00 2001
From: 3V3RYONE <beleswarprasad@gmail.com>
Date: Thu, 23 Jun 2022 16:50:23 +0530
Subject: [PATCH] check for NOEXEC flags before creating db socket file

---
 lib/msfdb_helpers/db_interface.rb |  6 ++---
 lib/msfdb_helpers/pg_ctl.rb       | 45 ++++++++++++++++++++++++-------
 msfdb                             |  1 +
 3 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/lib/msfdb_helpers/db_interface.rb b/lib/msfdb_helpers/db_interface.rb
index 5c826b049a..25d90a3e5c 100644
--- a/lib/msfdb_helpers/db_interface.rb
+++ b/lib/msfdb_helpers/db_interface.rb
@@ -60,12 +60,12 @@ module MsfdbHelpers
       status.exitstatus
     end
 
-    def run_psql(cmd, db_name: 'postgres')
+    def run_psql(cmd, socket_directory= "#{Dir.tmpdir}", db_name: 'postgres')
       if @options[:debug]
-        puts "psql -h #{Dir.tmpdir} -p #{@options[:db_port]} -c \"#{cmd};\" #{db_name}"
+        puts "psql -h #{socket_directory} -p #{@options[:db_port]} -c \"#{cmd};\" #{db_name}"
       end
 
-      run_cmd("psql -h #{Dir.tmpdir} -p #{@options[:db_port]} -c \"#{cmd};\" #{db_name}")
+      run_cmd("psql -h #{socket_directory} -p #{@options[:db_port]} -c \"#{cmd};\" #{db_name}")
     end
 
   end
diff --git a/lib/msfdb_helpers/pg_ctl.rb b/lib/msfdb_helpers/pg_ctl.rb
index b60f0b5e31..70ddd094c3 100644
--- a/lib/msfdb_helpers/pg_ctl.rb
+++ b/lib/msfdb_helpers/pg_ctl.rb
@@ -8,6 +8,7 @@ module MsfdbHelpers
       @options = options
       @localconf = localconf
       @db_conf = db_conf
+      @socket_directory = db_path
       super(options)
     end
 
@@ -18,10 +19,17 @@ module MsfdbHelpers
 
       File.open("#{@db}/postgresql.conf", 'a') do |f|
         f.puts "port = #{@options[:db_port]}"
-        if system("mount | grep #{Dir.tmpdir}.*noexec")
-          print_error("Temporary Directory is mounted with NOEXEC flags. Try running sudo msfdb init, if initialization fails")
-        else
-          f.puts "unix_socket_directories = \'#{Dir.tmpdir}\'"
+      end
+
+      # Try creating a test file at {Dir.tmpdir}
+      begin
+        test_executable_file("#{Dir.tmpdir}")
+      rescue
+        begin
+          # Fallback to creation at {@db}
+          test_executable_file("#{@db}")
+        rescue
+          print_error("Attempt to create DB socket file at Temporary Directory and `~/.msf4/db` failed. Possibly because they are mounted with NOEXEC Flags. Database Initialization Failed.")
         end
       end
 
@@ -33,6 +41,23 @@ module MsfdbHelpers
       restart
     end
 
+    # Creates and attempts to execute a testfile in the specified directory,
+    # to determine if it is mounted with NOEXEC flags.
+    def test_executable_file(path)
+      File.open("#{path}/msfdb_testfile", 'w') do |f|
+        f.puts "#!/bin/bash\necho exec"
+      end
+      File.chmod(0744, "#{path}/msfdb_testfile")
+      
+      if run_cmd("#{path}/msfdb_testfile")
+        File.open("#{@db}/postgresql.conf", 'a') do |f|
+          f.puts "unix_socket_directories = \'#{path}\'"
+        end
+        @socket_directory = path
+        puts "Creating db socket file at #{path}"
+      end
+    end
+
     def delete
       if exists?
         stop
@@ -100,12 +125,12 @@ module MsfdbHelpers
 
     def create_db_users(msf_pass, msftest_pass)
       puts 'Creating database users'
-      run_psql("create user #{@options[:msf_db_user].shellescape} with password '#{msf_pass}'")
-      run_psql("create user #{@options[:msftest_db_user].shellescape} with password '#{msftest_pass}'")
-      run_psql("alter role #{@options[:msf_db_user].shellescape} createdb")
-      run_psql("alter role #{@options[:msftest_db_user].shellescape} createdb")
-      run_psql("alter role #{@options[:msf_db_user].shellescape} with password '#{msf_pass}'")
-      run_psql("alter role #{@options[:msftest_db_user].shellescape} with password '#{msftest_pass}'")
+      run_psql("create user #{@options[:msf_db_user].shellescape} with password '#{msf_pass}'", @socket_directory)
+      run_psql("create user #{@options[:msftest_db_user].shellescape} with password '#{msftest_pass}'", @socket_directory)
+      run_psql("alter role #{@options[:msf_db_user].shellescape} createdb", @socket_directory)
+      run_psql("alter role #{@options[:msftest_db_user].shellescape} createdb", @socket_directory)
+      run_psql("alter role #{@options[:msf_db_user].shellescape} with password '#{msf_pass}'", @socket_directory)
+      run_psql("alter role #{@options[:msftest_db_user].shellescape} with password '#{msftest_pass}'", @socket_directory)
 
       conn = PG.connect(host: @options[:db_host], dbname: 'postgres', port: @options[:db_port], user: @options[:msf_db_user], password: msf_pass)
       conn.exec("CREATE DATABASE #{@options[:msf_db_name]}")
diff --git a/msfdb b/msfdb
index 1f6f3a3030..3b65c67558 100755
--- a/msfdb
+++ b/msfdb
@@ -205,6 +205,7 @@ def init_db
   Dir.chdir(@framework) do
     @db_driver.run_cmd('bundle exec rake db:migrate')
   end
+  puts 'Database initialization successful'.green.bold.to_s
 end
 
 def load_db_config