23s2s

✏️ Editing: source_finder.rb

# frozen_string_literal: true

require 'prism'

module IRB
  class SourceFinder
    class EvaluationError < StandardError; end

class Source
      attr_reader :file, :line
      def initialize(file, line, ast_source = nil)
        @file = file
        @line = line
        @ast_source = ast_source
      end

def file_exist?
        File.exist?(@file)
      end

def binary_file?
        # If the line is zero, it means that the target's source is probably in a binary file.
        @line.zero?
      end

def file_content
        @file_content ||= File.read(@file)
      end

def colorized_content
        if !binary_file? && file_exist?
          # To correctly colorize, we need to colorize full content and extract the relevant lines.
          colored_lines = IRB::Color.colorize_code(file_content).lines

# Handle wrong line number case: line_no passed to eval is wrong, file is edited after load, etc
          return if colored_lines.size < @line

colored_lines[@line - 1...find_end].join
        elsif @ast_source
          IRB::Color.colorize_code(@ast_source)
        end
      end

private

def find_end
        code = file_content
        lines = code.lines[(@line - 1)..-1]

lines.each_with_index do |line, index|
          sub_code = lines.take(index + 1).join
          return @line + index if Prism.parse_success?(sub_code)
        end
        @line
      end
    end

private_constant :Source

def initialize(irb_context)
      @irb_context = irb_context
    end

def find_source(signature, super_level = 0)
      case signature
      when /\A(::)?[A-Z]\w*(::[A-Z]\w*)*\z/ # ConstName, ::ConstName, ConstPath::ConstName
        eval_receiver_or_owner(signature) # trigger autoload
        *parts, name = signature.split('::', -1)
        base =
          if parts.empty? # ConstName
            find_const_owner(name)
          elsif parts == [''] # ::ConstName
            Object
          else # ConstPath::ConstName
            eval_receiver_or_owner(parts.join('::'))
          end
        file, line = base.const_source_location(name)
      when /\A(?<owner>[A-Z]\w*(::[A-Z]\w*)*)#(?<method>[^ :.]+)\z/ # Class#method
        owner = eval_receiver_or_owner(Regexp.last_match[:owner])
        method = Regexp.last_match[:method]
        return unless owner.respond_to?(:instance_method)
        method = method_target(owner, super_level, method, "owner")
        file, line = method&.source_location
      when /\A((?<receiver>.+)(\.|::))?(?<method>[^ :.]+)\z/ # method, receiver.method, receiver::method
        receiver = eval_receiver_or_owner(Regexp.last_match[:receiver] || 'self')
        method = Regexp.last_match[:method]
        return unless receiver.respond_to?(method, true)
        method = method_target(receiver, super_level, method, "receiver")
        file, line = method&.source_location
      end
      return unless file && line

if File.exist?(file)
        Source.new(file, line)
      elsif method
        # Method defined with eval, probably in IRB session
        source = RubyVM::InstructionSequence.of(method)&.script_lines&.join rescue nil
        Source.new(file, line, source)
      end
    rescue EvaluationError
      nil
    end

private

def method_target(owner_receiver, super_level, method, type)
      case type
      when "owner"
        target_method = owner_receiver.instance_method(method)
      when "receiver"
        target_method = Kernel.instance_method(:method).bind_call(owner_receiver, method)
      end
      super_level.times do |s|
        target_method = target_method.super_method if target_method
      end
      target_method
    rescue NameError
      nil
    end

def eval_receiver_or_owner(code)
      @irb_context.workspace.binding.eval(code)
    rescue Exception
      raise EvaluationError
    end

def find_const_owner(name)
      module_nesting = @irb_context.workspace.binding.eval('::Module.nesting')
      module_nesting.find { |mod| mod.const_defined?(name, false) } || module_nesting.find { |mod| mod.const_defined?(name) } || Object
    end
  end
end

Sid Gifari From Gifari Industries - BD Cyber Security Team