Class: Omamori::CoreRunner

Inherits:
Object
  • Object
show all
Defined in:
lib/omamori/core_runner.rb

Constant Summary collapse

JSON_SCHEMA =

Define the JSON Schema for Structured Output

{
  "type": "object",
  "properties": {
    "security_risks": {
      "type": "array",
      "description": "検出されたセキュリティリスクのリスト。",
      "items": {
        "type": "object",
        "properties": {
          "type": {
            "type": "string",
            "description": "検出されたリスクの種類 (例: XSS, CSRF, IDORなど)。3.3の診断対象脆弱性リストのいずれかであること。"
          },
          "location": {
            "type": "string",
            "description": "リスクが存在するコードのファイル名、行番号、またはコードスニペット。差分分析の場合は差分の該当箇所を示す形式 (例: ファイル名:+行番号) であること。"
          },
          "details": {
            "type": "string",
            "description": "リスクの詳細な説明と、なぜそれがリスクなのかの理由。"
          },
          "severity": {
            "type": "string",
            "description": "リスクの深刻度。",
            "enum": ["Critical", "High", "Medium", "Low", "Info"]
          },
          "code_snippet": {
            "type": "string",
            "description": "該当するコードスニペット。"
          }
        },
        "required": ["type", "location", "details", "severity"]
      }
    }
  },
  "required": ["security_risks"]
}.freeze
RISKS_TO_CHECK =

TODO: Get risks to check from config file

[
  :xss, :csrf, :idor, :open_redirect, :ssrf, :session_fixation
  # TODO: Add other risks from requirements
].freeze
SPLIT_THRESHOLD =

TODO: Determine threshold for splitting based on token limits

7000

Instance Method Summary collapse

Constructor Details

#initialize(args) ⇒ CoreRunner

Characters as a proxy for tokens



65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
 
# File 'lib/omamori/core_runner.rb', line 65

def initialize(args)
  @args = args
  @options = { command: :scan, format: :console } # Default command is scan, default format is console
  @config = Omamori::Config.new # Initialize Config

  # Initialize components with config
  api_key = @config.get("api_key", ENV["GEMINI_API_KEY"]) # Get API key from config or environment variable
  gemini_model = @config.get("model", "gemini-1.5-pro-latest") # Get Gemini model from config
  @gemini_client = AIAnalysisEngine::GeminiClient.new(api_key)
  @prompt_manager = AIAnalysisEngine::PromptManager.new(@config) # Pass the entire config object
  # Get chunk size from config, default to 7000 characters if not specified
  chunk_size = @config.get("chunk_size", SPLIT_THRESHOLD)
  @diff_splitter = AIAnalysisEngine::DiffSplitter.new(chunk_size: chunk_size) # Pass chunk size to DiffSplitter
  # Get report output path and html template path from config
  report_config = @config.get("report", {})
  report_output_path = report_config.fetch("output_path", "./omamori_report")
  html_template_path = report_config.fetch("html_template", nil) # Default to nil, formatter will use default template
  @console_formatter = ReportGenerator::ConsoleFormatter.new # TODO: Pass config for colors/options
  @html_formatter = ReportGenerator::HTMLFormatter.new(report_output_path, html_template_path) # Pass output path and template path
  @json_formatter = ReportGenerator::JSONFormatter.new(report_output_path) # Pass output path
  # Get static analyser options from config
  static_analyser_config = @config.get("static_analysers", {})
  brakeman_options = static_analyser_config.fetch("brakeman", {}).fetch("options", {}) # Default to empty hash
  bundler_audit_options = static_analyser_config.fetch("bundler_audit", {}).fetch("options", {}) # Default to empty hash
  @brakeman_runner = StaticAnalysers::BrakemanRunner.new(brakeman_options) # Pass options
  @bundler_audit_runner = StaticAnalysers::BundlerAuditRunner.new(bundler_audit_options) # Pass options
end

Instance Method Details

#run ⇒ Object 



93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
 
# File 'lib/omamori/core_runner.rb', line 93

def run
  parse_options

  case @options[:command]
  when :scan
    # Run static analysers first unless --ai option is specified
    brakeman_result = nil
    bundler_audit_result = nil
    unless @options[:only_ai]
      brakeman_result = @brakeman_runner.run
      bundler_audit_result = @bundler_audit_runner.run
    end

    # Perform AI analysis
    analysis_result = nil
    case @options[:scan_mode]
    when :diff
      diff_content = get_staged_diff
      if diff_content.empty?
        puts "No staged changes to scan."
        return
      end
      puts "Scanning staged differences with AI..."
      if diff_content.length > SPLIT_THRESHOLD # TODO: Use token count
        puts "Diff content exceeds threshold, splitting..."
        analysis_result = @diff_splitter.process_in_chunks(diff_content, @gemini_client, JSON_SCHEMA, @prompt_manager, get_risks_to_check, model: @config.get("model", "gemini-1.5-pro-latest"))
      else
        prompt = @prompt_manager.build_prompt(diff_content, get_risks_to_check, JSON_SCHEMA)
        analysis_result = @gemini_client.analyze(prompt, JSON_SCHEMA, model: @config.get("model", "gemini-1.5-pro-latest"))
      end
    when :all
      full_code_content = get_full_codebase
      if full_code_content.strip.empty?
        puts "No code found to scan."
        return
      end
      puts "Scanning entire codebase with AI..."
      if full_code_content.length > SPLIT_THRESHOLD # TODO: Use token count
        puts "Full code content exceeds threshold, splitting..."
        analysis_result = @diff_splitter.process_in_chunks(full_code_content, @gemini_client, JSON_SCHEMA, @prompt_manager, get_risks_to_check, model: @config.get("model", "gemini-1.5-pro-latest"))
      else
        prompt = @prompt_manager.build_prompt(full_code_content, get_risks_to_check, JSON_SCHEMA)
        analysis_result = @gemini_client.analyze(prompt, JSON_SCHEMA, model: @config.get("model", "gemini-1.5-pro-latest"))
      end
    end

    # Combine results and display report
    combined_results = combine_results(analysis_result, brakeman_result, bundler_audit_result)
    display_report(combined_results)

    puts "Scan complete."

  when :ci_setup
    generate_ci_setup(@options[:ci_service])

  when :init
    generate_config_file # Generate initial config file

  else
    puts "Unknown command: #{@options[:command]}"
    puts @opt_parser # Display help for unknown command
  end
end