From 83d6c146259a612e892e0714706085f86220562f Mon Sep 17 00:00:00 2001
From: Bastian Schmidt <schmidt@atix.de>
Date: Thu, 14 Jul 2022 15:33:47 +0200
Subject: [PATCH] Enable boot image download for iso images

* Implement fetch and extract boot image
* Apply correct file permissions
* Introduce and add tests for fetch_boot_image
* Implement class for file extraction
---
 lib/proxy/archive_extract.rb | 30 ++++++++++++++++++++++++++++++
 lib/proxy/util.rb            | 12 +++++++++---
 lib/smart_proxy_main.rb      |  1 +
 modules/tftp/server.rb       | 20 ++++++++++++++++++++
 modules/tftp/tftp_api.rb     |  4 ++++
 test/tftp/tftp_api_test.rb   |  8 +++++++-
 6 files changed, 71 insertions(+), 4 deletions(-)
 create mode 100644 lib/proxy/archive_extract.rb

diff --git a/lib/proxy/archive_extract.rb b/lib/proxy/archive_extract.rb
new file mode 100644
index 000000000..ec128c2c5
--- /dev/null
+++ b/lib/proxy/archive_extract.rb
@@ -0,0 +1,30 @@
+module Proxy
+  class ArchiveExtract < Proxy::Util::CommandTask
+    include Util
+
+    def initialize(image_path, file_in_image, dst_path)
+      args = [which('isoinfo')]
+
+      # read the file
+      args << "-R"
+      # set image path
+      args += ["-i", image_path.to_s]
+      # set file path within the image
+      args += ["-x", file_in_image.to_s]
+
+      super(args, input=nil, output=dst_path)
+    end
+
+    def start
+      lock = Proxy::FileLock.try_locking(File.join(File.dirname(@output), ".#{File.basename(@output)}.lock"))
+      if lock.nil?
+        false
+      else
+        super do
+          Proxy::FileLock.unlock(lock)
+          File.unlink(lock)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/proxy/util.rb b/lib/proxy/util.rb
index e191a70dd..ed4302471 100644
--- a/lib/proxy/util.rb
+++ b/lib/proxy/util.rb
@@ -13,9 +13,10 @@ class CommandTask
     # stderr is redirected to proxy error log, stdout to proxy debug log
     # command can be either string or array (command + arguments)
     # input is passed into STDIN and must be string
-    def initialize(command, input = nil)
+    def initialize(command, input = nil, output = nil)
       @command = command
       @input = input
+      @output = output
     end
 
     def start(&ensured_block)
@@ -27,8 +28,13 @@ def start(&ensured_block)
           logger.info "[#{thr.pid}] Started task #{cmdline_string}"
           stdin.write(input) if input
           stdin.close
-          stdout.each do |line|
-            logger.debug "[#{thr.pid}] #{line}"
+          unless @output.nil?
+            File.binwrite(@output, stdout.read)
+            stdout.close
+          else
+            stdout.each do |line|
+              logger.debug "[#{thr.pid}] #{line}"
+            end
           end
           stderr.each do |line|
             logger.warn "[#{thr.pid}] #{line}"
diff --git a/lib/smart_proxy_main.rb b/lib/smart_proxy_main.rb
index 0ebc22b2c..cb1217cf4 100644
--- a/lib/smart_proxy_main.rb
+++ b/lib/smart_proxy_main.rb
@@ -14,6 +14,7 @@
 require 'proxy/dependency_injection'
 require 'proxy/util'
 require 'proxy/http_download'
+require 'proxy/archive_extract'
 require 'proxy/helpers'
 require 'proxy/memory_store'
 require 'proxy/plugin_validators'
diff --git a/modules/tftp/server.rb b/modules/tftp/server.rb
index eba4e5e4a..79a4b6bd2 100644
--- a/modules/tftp/server.rb
+++ b/modules/tftp/server.rb
@@ -150,6 +150,26 @@ def pxeconfig_file(mac)
     end
   end
 
+  def self.fetch_boot_image(image_dst, url, files)
+    # Verify dst is a valid directory
+    image_path = Pathname.new(image_dst).cleanpath
+    extr_image_dir = Pathname.new(image_dst.delete_suffix(".iso"))
+
+    FileUtils.mkdir_p image_path.parent
+    choose_protocol_and_fetch(url, image_path).join
+
+    files.each do |file|
+      file_path = Pathname.new file
+      extr_file_path = Pathname.new(File.join(extr_image_dir, file_path)).cleanpath
+
+      # Create destination directory
+      FileUtils.mkdir_p extr_file_path.parent
+      # extract iso
+      extract_task = ::Proxy::ArchiveExtract.new(image_path, file_path, extr_file_path).start
+      raise "TFTP image file extraction error: #{file_path}" unless extract_task.join == 0
+    end
+  end
+
   def self.fetch_boot_file(dst, src)
     filename    = boot_filename(dst, src)
     destination = Pathname.new(File.expand_path(filename, Proxy::TFTP::Plugin.settings.tftproot)).cleanpath
diff --git a/modules/tftp/tftp_api.rb b/modules/tftp/tftp_api.rb
index 1bc6276c7..4f8f2c4da 100644
--- a/modules/tftp/tftp_api.rb
+++ b/modules/tftp/tftp_api.rb
@@ -34,6 +34,10 @@ def create_default(variant)
       end
     end
 
+    post "/fetch_boot_image" do
+      log_halt(400, "TFTP: Failed to fetch boot file: ") { Proxy::TFTP.fetch_boot_image(params[:path], params[:url], params[:files]) }
+    end
+
     post "/fetch_boot_file" do
       log_halt(400, "TFTP: Failed to fetch boot file: ") { Proxy::TFTP.fetch_boot_file(params[:prefix], params[:path]) }
     end
diff --git a/test/tftp/tftp_api_test.rb b/test/tftp/tftp_api_test.rb
index 3ffcfb0cb..9e9270ada 100644
--- a/test/tftp/tftp_api_test.rb
+++ b/test/tftp/tftp_api_test.rb
@@ -111,7 +111,13 @@ def test_api_can_fetch_boot_file
     assert last_response.ok?
   end
 
-  def test_api_can_get_servername
+  def test_api_can_fetch_boot_image
+    Proxy::TFTP.expects(:fetch_boot_image).with('some/image.iso', 'http://localhost/file.iso').returns(true)
+    post "/fetch_boot_image", :path => 'some/image.iso', :url => 'http://localhost/file.iso'
+    assert last_response.ok?
+  end
+
+ def test_api_can_get_servername
     Proxy::TFTP::Plugin.settings.stubs(:tftp_servername).returns("servername")
     result = get "/serverName"
     assert_match /servername/, result.body