<commit_before>module AWSSecurityGroups
  class Updater
    def initialize(args)
      @product, @region, @owner_id = args[:product], args[:region], args[:owner_id]

      if @product == "ec2"
        @aws = Fog::Compute.new(:provider => "AWS", :region => @region, :aws_access_key_id => args[:access_key], :aws_secret_access_key => args[:secret_key])
      elsif @product == "rds"
        @aws = Fog::AWS::RDS.new(:region => @region, :aws_access_key_id => args[:access_key], :aws_secret_access_key => args[:secret_key])
      end
    end

    def add_rules(group_name, rules)
      self.send("add_#{@product}_rules", group_name, rules)
    end

    def remove_rules(group_name, rules)
      self.send("remove_#{@product}_rules", group_name, rules)
    end

    private
    # EC2
    def add_ec2_rules(group_name, rules)
      rules.each do |rule|
        if rule[:ip]
          data = {"CidrIp" => rule[:ip], "IpProtocol" => rule[:protocol], "FromPort" => rule[:port], "ToPort" => rule[:port]}
        else
          data = {"SourceSecurityGroupName" => rule[:group], "SourceSecurityGroupOwnerId" => @owner_id}
        end

        puts "AUTHORIZE: #{data}"
        @aws.authorize_security_group_ingress(group_name, data)
      end
    end

    def remove_ec2_rules(group_name, rules)
      rules.each do |rule|
        if rule[:ip]
          data = {"CidrIp" => rule[:ip], "IpProtocol" => rule[:protocol], "FromPort" => rule[:port], "ToPort" => rule[:port]}
        else
          data = {"SourceSecurityGroupName" => rule[:group], "SourceSecurityGroupOwnerId" => @owner_id}
        end

        puts "REVOKING: #{data}"
        @aws.revoke_security_group_ingress(group_name, data)
      end
    end

    # RDS
    def add_rds_rules(group_name, rules)
      rules.each do |rule|
        if rule[:ip]
          data = {"CIDRIP" => rule[:ip]}
        else
          data = {"EC2SecurityGroupName" => rule[:group], "EC2SecurityGroupOwnerId" => @owner_id}
        end

        puts "AUTHORIZE: #{data}"
        @aws.authorize_db_security_group_ingress_ingress(group_name, {"CIDRIP" => rule[:ip]})
      end
    end

    def remove_rds_rules(group_name, rules)
      rules.each do |rule|
        if rule[:ip]
          data = {"CIDRIP" => rule[:ip]}
        else
          data = {"EC2SecurityGroupName" => rule[:group], "EC2SecurityGroupOwnerId" => @owner_id}
        end

        puts "REVOKING: #{data}"
        @aws.revoke_db_security_group_ingress(group_name, data)
      end
    end
  end
end<commit_msg>Added support for port ranges<commit_after>module AWSSecurityGroups
  class Updater
    def initialize(args)
      @product, @region, @owner_id = args[:product], args[:region], args[:owner_id]

      if @product == "ec2"
        @aws = Fog::Compute.new(:provider => "AWS", :region => @region, :aws_access_key_id => args[:access_key], :aws_secret_access_key => args[:secret_key])
      elsif @product == "rds"
        @aws = Fog::AWS::RDS.new(:region => @region, :aws_access_key_id => args[:access_key], :aws_secret_access_key => args[:secret_key])
      end
    end

    def add_rules(group_name, rules)
      self.send("add_#{@product}_rules", group_name, rules)
    end

    def remove_rules(group_name, rules)
      self.send("remove_#{@product}_rules", group_name, rules)
    end

    private
    # EC2
    def format_ec2_rule(rule)
      if rule[:ip]
        if rule[:port] =~ /\-/
          from, to = rule[:port].split("-", 2)
        else
          from, to = rule[:port], rule[:port]
        end

        {"CidrIp" => rule[:ip], "IpProtocol" => rule[:protocol], "FromPort" => from.to_i, "ToPort" => to.to_i}
      else
        {"SourceSecurityGroupName" => rule[:group], "SourceSecurityGroupOwnerId" => @owner_id}
      end
    end

    def add_ec2_rules(group_name, rules)
      rules.each do |rule|
        data = format_ec2_rule(rule)

        puts "AUTHORIZE: #{data}"
        @aws.authorize_security_group_ingress(group_name, data)
      end
    end

    def remove_ec2_rules(group_name, rules)
      rules.each do |rule|
        data = format_ec2_rule(rule)

        puts "REVOKING: #{data}"
        @aws.revoke_security_group_ingress(group_name, data)
      end
    end

    # RDS
    def add_rds_rules(group_name, rules)
      rules.each do |rule|
        if rule[:ip]
          data = {"CIDRIP" => rule[:ip]}
        else
          data = {"EC2SecurityGroupName" => rule[:group], "EC2SecurityGroupOwnerId" => @owner_id}
        end

        puts "AUTHORIZE: #{data}"
        @aws.authorize_db_security_group_ingress_ingress(group_name, {"CIDRIP" => rule[:ip]})
      end
    end

    def remove_rds_rules(group_name, rules)
      rules.each do |rule|
        if rule[:ip]
          data = {"CIDRIP" => rule[:ip]}
        else
          data = {"EC2SecurityGroupName" => rule[:group], "EC2SecurityGroupOwnerId" => @owner_id}
        end

        puts "REVOKING: #{data}"
        @aws.revoke_db_security_group_ingress(group_name, data)
      end
    end
  end
end
{'fraction_non_alphanumeric': 0.14062186559679038, 'fraction_numerical': 0.004012036108324975, 'mean_word_length': 2.8061068702290077, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 49, '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': '733233', 'n_tokens_mistral': 1684, 'n_tokens_neox': 1675, 'n_words': 338}