<commit_before>require 'aws-sdk'
require 'yaml'

module Ec2SecurityCzar
  class Base
    attr_accessor :ec2

    def initialize(mfa_token=nil)
      keys = YAML.load_file('config/aws_keys.yml')
      AWS.config(access_key_id: keys[:access_key], secret_access_key: keys[:secret_key], region: "us-east-1")
      if keys[:mfa_serial_number]
        @ec2 = mfa_auth(keys, mfa_token)
      else
        @ec2 = AWS.ec2
      end
    rescue StandardError => e
      puts e.class
      puts e.message
      exit 1
    end

    def update_rules
      security_groups.each do |sg|
        security_group = SecurityGroup.new sg
        security_group.update_rules
      end
    end

    def security_groups
      ec2.security_groups
    end

    private
    def mfa_auth(keys, mfa_token)
      raise MFATokenMissing unless mfa_token
      sts = AWS::STS.new(access_key_id: keys[:access_key], secret_access_key: keys[:secret_key])
      session = sts.new_session(duration: keys[:mfa_duration] || 900, serial_number: keys[:mfa_serial_number], token_code: mfa_token)
      AWS::EC2.new(session.credentials)
    rescue MFATokenMissing
      puts "MFA token is required as an argument!"
      exit 1
    rescue AWS::STS::Errors::AccessDenied
      puts "Multifactor Authentication failed."
      exit 1
    end
  end

  class MFATokenMissing < StandardError
  end
end<commit_msg>change error handling<commit_after>require 'aws-sdk'
require 'yaml'

module Ec2SecurityCzar
  class Base
    attr_accessor :ec2

    def initialize(mfa_token=nil)
      keys = YAML.load_file('config/aws_keys.yml')
      AWS.config(access_key_id: keys[:access_key], secret_access_key: keys[:secret_key], region: "us-east-1")
      if keys[:mfa_serial_number]
        @ec2 = mfa_auth(keys, mfa_token)
      else
        @ec2 = AWS.ec2
      end
    rescue StandardError => e
      case
      when e.class == Ec2SecurityCzar::MFATokenMissing
        puts "MFA token is required as an argument!"
      else
        puts e.class
        puts e.message
      end
      exit 1
    end

    def update_rules
      security_groups.each do |sg|
        security_group = SecurityGroup.new sg
        security_group.update_rules
      end
    end

    def security_groups
      ec2.security_groups
    end

    private
    def mfa_auth(keys, mfa_token)
      raise MFATokenMissing unless mfa_token
      sts = AWS::STS.new(access_key_id: keys[:access_key], secret_access_key: keys[:secret_key])
      session = sts.new_session(duration: keys[:mfa_duration] || 900, serial_number: keys[:mfa_serial_number], token_code: mfa_token)
      AWS::EC2.new(session.credentials)
    end
  end

  MFATokenMissing = Class.new StandardError
end
{'fraction_non_alphanumeric': 0.08389261744966443, 'fraction_numerical': 0.010067114093959731, 'mean_word_length': 2.922514619883041, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '217645', 'n_tokens_mistral': 903, 'n_tokens_neox': 878, 'n_words': 200}