From a3b71ad0510070b3d86625e484b52a0a0bf8512e Mon Sep 17 00:00:00 2001 From: elg0nz Date: Fri, 27 Sep 2024 13:47:41 -0700 Subject: [PATCH 1/2] add templating capabilities --- lib/request_interceptor.rb | 109 ++++++++++++++++++++++++++++++------- lib/uploads_cli.rb | 39 +++++++++++++ 2 files changed, 128 insertions(+), 20 deletions(-) diff --git a/lib/request_interceptor.rb b/lib/request_interceptor.rb index 4bb7702..8475647 100644 --- a/lib/request_interceptor.rb +++ b/lib/request_interceptor.rb @@ -1,5 +1,6 @@ # frozen_string_literal: true +require 'socket' require 'json' require 'logger' require 'digest' @@ -17,47 +18,114 @@ def initialize(app, logger: Logger.new('hawksi.log'), storage: FileStorage) @app = app @logger = logger @storage = storage + @templates = {} + @templates_mutex = Mutex.new + + @socket_path = '/tmp/hawksi.sock' + start_unix_socket_server end def call(env) request = Rack::Request.new(env) + request_key = generate_request_key(request) # Generate a key for the request + + # Check if there's a template for this request + template = nil + @templates_mutex.synchronize do + template = @templates[request_key] + end + if template + status, headers, _temp = @app.call(env) + + # Serve the template + headers['Etag'] = Digest::MD5.hexdigest(template) + return [status, headers, [template]] + end + # Original code return MocksiHandler.handle(request) if request.path.end_with?('/favicon.ico') if request.path.start_with?('/mocksi') || request.path.start_with?('/_') || request.path.start_with?('/api') return MocksiHandler.handle(request) end - request_hash = generate_request_hash(request) # Generate a hash of the request - log_request(request, request_hash) + log_request(request, request_key) status, headers, response = @app.call(env) - log_response(status, headers, response, request_hash) + log_response(status, headers, response, request_key) [status, headers, response] end private - def generate_request_hash(request) - # Generate a hash based on request method, path, query string, and body - hash_input = [ - request.request_method, - request.path, - request.query_string, - request.body&.read # Read the body content to include in the hash - ].join + def start_unix_socket_server + # Remove the socket file if it already exists + File.unlink(@socket_path) if File.exist?(@socket_path) - # Reset the body input stream for future use - request.body&.rewind + at_exit do + File.unlink(@socket_path) if File.exist?(@socket_path) + end + + @server_thread = Thread.new do + @server = UNIXServer.new(@socket_path) + @logger.info("Unix socket server started at #{@socket_path}") + puts "Unix socket server started at #{@socket_path}" + loop do + client = @server.accept + Thread.new(client) do |conn| + handle_client(conn) + end + end + rescue StandardError => e + @logger.error("Unix socket server error: #{e.message}") + ensure + @server.close if @server + end + end - # Return a SHA256 hash of the concatenated string - Digest::SHA256.hexdigest(hash_input) + def handle_client(conn) + while message = conn.gets + # Process the message + process_message(message.chomp) + end + rescue StandardError => e + @logger.error("Error handling client: #{e.message}") + ensure + conn.close end - def log_request(request, request_hash) # rubocop:disable Metrics/AbcSize,Metrics/MethodLength + def process_message(message) + data = JSON.parse(message) + case data['action'] + when 'set_template' + key = data['key'] + template = data['template'] + @templates_mutex.synchronize do + @templates[key] = template + end + @logger.info("Template set for key: #{key}") + when 'remove_template' + key = data['key'] + @templates_mutex.synchronize do + @templates.delete(key) + end + @logger.info("Template removed for key: #{key}") + else + @logger.warn("Unknown action: #{data['action']}") + end + rescue JSON::ParserError => e + puts("JSON parsing error: #{e.message}") + @logger.error("JSON parsing error: #{e.message}") + end + + def generate_request_key(request) + "#{request.request_method}_#{request.path}" + end + + def log_request(request, request_key) # rubocop:disable Metrics/AbcSize,Metrics/MethodLength data = { - request_hash: request_hash, # Include the request hash in the logged data + request_key: request_key, # Include the request key in the logged data method: request.request_method, path: request.path, query_string: request.query_string, @@ -65,7 +133,6 @@ def log_request(request, request_hash) # rubocop:disable Metrics/AbcSize,Metrics scheme: request.scheme, host: request.host, port: request.port, - # Log only specific parts of the env hash to avoid circular references env: { 'REQUEST_METHOD' => request.env['REQUEST_METHOD'], 'SCRIPT_NAME' => request.env['SCRIPT_NAME'], @@ -95,16 +162,18 @@ def log_request(request, request_hash) # rubocop:disable Metrics/AbcSize,Metrics @storage.store('requests', data) rescue StandardError => e @logger.error("Error logging request: #{e.message}") + ensure + request.body&.rewind end - def log_response(status, headers, response, request_hash) # rubocop:disable Metrics/MethodLength + def log_response(status, headers, response, request_key) # rubocop:disable Metrics/MethodLength body = if response.respond_to?(:body) response.body.join.to_s else response.join.to_s end data = { - request_hash: request_hash, # Include the request hash in the response log + request_key: request_key, # Include the request key in the response log status: status, headers: headers, body: body, diff --git a/lib/uploads_cli.rb b/lib/uploads_cli.rb index 90ae362..8bcc5ce 100644 --- a/lib/uploads_cli.rb +++ b/lib/uploads_cli.rb @@ -26,6 +26,7 @@ def initialize(*args) @client_uuid = current_client_uuid @file_uploader = FileUploader.new(@logger, @client_uuid) @command_executor = CommandExecutor.new(@logger, @client_uuid) + @socket_path = '/tmp/hawksi.sock' end desc 'update', 'Update uploaded requests and responses' @@ -62,6 +63,44 @@ def execute(command, *params) @command_executor.execute_command(command, params) end + desc 'set_template KEY FILENAME', 'Set a template for the given key using the content of the specified file' + option :base_dir, type: :string, + desc: 'Base directory for storing intercepted data. Defaults to ./tmp/intercepted_data' + def set_template(key, filename) + set_base_dir + unless File.exist?(filename) + @logger.error "File not found: #{filename}" + return + end + + # Read the template content from the file + template_content = File.read(filename) + + message = { + action: 'set_template', + key: key, + template: template_content + } + # Send the message to the /tmp/hawksi.sock socket + socket = UNIXSocket.new(@socket_path) + socket.write(message.to_json) + socket.close + end + + desc 'remove_template KEY', 'Remove the template for the given key' + option :base_dir, type: :string, + desc: 'Base directory for storing intercepted data. Defaults to ./tmp/intercepted_data' + def remove_template(key) + set_base_dir + message = { + action: 'remove_template', + key: key + } + socket = UNIXSocket.new(@socket_path) + socket.write(message.to_json) + socket.close + end + private def set_base_dir From 965803113826cb8538ab5e49a251101471542e46 Mon Sep 17 00:00:00 2001 From: elg0nz Date: Mon, 30 Sep 2024 19:18:28 -0700 Subject: [PATCH 2/2] log result to stdout for demo purposes --- lib/command_executor.rb | 1 + 1 file changed, 1 insertion(+) diff --git a/lib/command_executor.rb b/lib/command_executor.rb index 135cd54..9cb8c75 100644 --- a/lib/command_executor.rb +++ b/lib/command_executor.rb @@ -61,6 +61,7 @@ def log_command_result(result) logger.error "Error during command execution: #{result['message']}" else logger.info "Command executed successfully. #{result}" + puts(result.fetch('response', '')) end end end