<commit_before>require 'timeout'

require_relative 'gitlab_config'
require_relative 'gitlab_logger'
require_relative 'gitlab_metrics'

class GitlabKeys # rubocop:disable Metrics/ClassLength
  class KeyError < StandardError; end

  attr_accessor :auth_file, :key

  def self.command(whatever)
    "#{ROOT_PATH}/bin/gitlab-shell #{whatever}"
  end

  def self.command_key(key_id)
    unless /\A[a-z0-9-]+\z/ =~ key_id
      raise KeyError, "Invalid key_id: #{key_id.inspect}"
    end

    command(key_id)
  end

  def self.whatever_line(command, trailer)
    "command=\"#{command}\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty #{trailer}"
  end

  def self.key_line(key_id, public_key)
    public_key.chomp!

    if public_key.include?("\n")
      raise KeyError, "Invalid public_key: #{public_key.inspect}"
    end

    whatever_line(command_key(key_id), public_key)
  end

  def self.principal_line(username_key_id, principal)
    principal.chomp!

    if principal.include?("\n")
      raise KeyError, "Invalid principal: #{principal.inspect}"
    end

    whatever_line(command_key(username_key_id), principal)
  end

  def initialize
    @command = ARGV.shift
    @key_id = ARGV.shift
    key = ARGV.shift
    @key = key.dup if key
    @auth_file = GitlabConfig.new.auth_file
  end

  def exec
    GitlabMetrics.measure("command-#{@command}") do
      case @command
      when 'add-key'
        add_key
      when 'batch-add-keys'
        batch_add_keys
      when 'rm-key'
        rm_key
      when 'list-keys'
        list_keys
      when 'list-key-ids'
        list_key_ids
      when 'clear'
        clear
      when 'check-permissions'
        check_permissions
      else
        $logger.warn('Attempt to execute invalid gitlab-keys command', command: @command.inspect)
        puts 'not allowed'
        false
      end
    end
  end

  protected

  def add_key
    lock do
      $logger.info('Adding key', key_id: @key_id, public_key: @key)
      auth_line = self.class.key_line(@key_id, @key)
      open_auth_file('a') { |file| file.puts(auth_line) }
    end
    true
  end

  def list_keys
    $logger.info 'Listing all keys'
    keys = ''
    File.readlines(auth_file).each do |line|
      # key_id & public_key
      # command=".../bin/gitlab-shell key-741" ... ssh-rsa AAAAB3NzaDAxx2E\n
      #                               ^^^^^^^              ^^^^^^^^^^^^^^^
      matches = /^command=\".+?\s+(.+?)\".+?(?:ssh|ecdsa)-.*?\s(.+)\s*.*\n*$/.match(line)
      keys << "#{matches[1]} #{matches[2]}\n" unless matches.nil?
    end
    keys
  end

  def list_key_ids
    $logger.info 'Listing all key IDs'
    open_auth_file('r') do |f|
      f.each_line do |line|
        matchd = line.match(/key-(\d+)/)
        next unless matchd
        puts matchd[1]
      end
    end
  end

  def batch_add_keys
    lock(300) do # Allow 300 seconds (5 minutes) for batch_add_keys
      open_auth_file('a') do |file|
        stdin.each_line do |input|
          tokens = input.strip.split("\t")
          abort("#{$0}: invalid input #{input.inspect}") unless tokens.count == 2
          key_id, public_key = tokens
          $logger.info('Adding key', key_id: key_id, public_key: public_key)
          file.puts(self.class.key_line(key_id, public_key))
        end
      end
    end
    true
  end

  def stdin
    $stdin
  end

  def rm_key
    lock do
      $logger.info('Removing key', key_id: @key_id)
      open_auth_file('r+') do |f|
        while line = f.gets # rubocop:disable Style/AssignmentInCondition
          next unless line.start_with?("command=\"#{self.class.command_key(@key_id)}\"")
          f.seek(-line.length, IO::SEEK_CUR)
          # Overwrite the line with #'s. Because the 'line' variable contains
          # a terminating '\n', we write line.length - 1 '#' characters.
          f.write('#' * (line.length - 1))
        end<commit_msg>Fix a RuboCop warning<commit_after>require 'timeout'

require_relative 'gitlab_config'
require_relative 'gitlab_logger'
require_relative 'gitlab_metrics'

class GitlabKeys # rubocop:disable Metrics/ClassLength
  class KeyError < StandardError; end

  attr_accessor :auth_file, :key

  def self.command(whatever)
    "#{ROOT_PATH}/bin/gitlab-shell #{whatever}"
  end

  def self.command_key(key_id)
    unless /\A[a-z0-9-]+\z/ =~ key_id
      raise KeyError, "Invalid key_id: #{key_id.inspect}"
    end

    command(key_id)
  end

  def self.whatever_line(command, trailer)
    "command=\"#{command}\",no-port-forwarding,no-X11-forwarding,no-agent-forwarding,no-pty #{trailer}"
  end

  def self.key_line(key_id, public_key)
    public_key.chomp!

    if public_key.include?("\n")
      raise KeyError, "Invalid public_key: #{public_key.inspect}"
    end

    whatever_line(command_key(key_id), public_key)
  end

  def self.principal_line(username_key_id, principal)
    principal.chomp!

    if principal.include?("\n")
      raise KeyError, "Invalid principal: #{principal.inspect}"
    end

    whatever_line(command_key(username_key_id), principal)
  end

  def initialize
    @command = ARGV.shift
    @key_id = ARGV.shift
    key = ARGV.shift
    @key = key.dup if key
    @auth_file = GitlabConfig.new.auth_file
  end

  def exec
    GitlabMetrics.measure("command-#{@command}") do
      case @command
      when 'add-key'
        add_key
      when 'batch-add-keys'
        batch_add_keys
      when 'rm-key'
        rm_key
      when 'list-keys'
        list_keys
      when 'list-key-ids'
        list_key_ids
      when 'clear'
        clear
      when 'check-permissions'
        check_permissions
      else
        $logger.warn('Attempt to execute invalid gitlab-keys command', command: @command.inspect)
        puts 'not allowed'
        false
      end
    end
  end

  protected

  def add_key
    lock do
      $logger.info('Adding key', key_id: @key_id, public_key: @key)
      auth_line = self.class.key_line(@key_id, @key)
      open_auth_file('a') { |file| file.puts(auth_line) }
    end
    true
  end

  def list_keys
    $logger.info 'Listing all keys'
    keys = ''
    File.readlines(auth_file).each do |line|
      # key_id & public_key
      # command=".../bin/gitlab-shell key-741" ... ssh-rsa AAAAB3NzaDAxx2E\n
      #                               ^^^^^^^              ^^^^^^^^^^^^^^^
      matches = /^command=\".+?\s+(.+?)\".+?(?:ssh|ecdsa)-.*?\s(.+)\s*.*\n*$/.match(line)
      keys << "#{matches[1]} #{matches[2]}\n" unless matches.nil?
    end
    keys
  end

  def list_key_ids
    $logger.info 'Listing all key IDs'
    open_auth_file('r') do |f|
      f.each_line do |line|
        matchd = line.match(/key-(\d+)/)
        next unless matchd
        puts matchd[1]
      end
    end
  end

  def batch_add_keys
    lock(300) do # Allow 300 seconds (5 minutes) for batch_add_keys
      open_auth_file('a') do |file|
        stdin.each_line do |input|
          tokens = input.strip.split("\t")
          abort("#{$0}: invalid input #{input.inspect}") unless tokens.count == 2
          key_id, public_key = tokens
          $logger.info('Adding key', key_id: key_id, public_key: public_key)
          file.puts(self.class.key_line(key_id, public_key))
        end
      end
    end
    true
  end

  def stdin
    $stdin
  end

  def rm_key
    lock do
      $logger.info('Removing key', key_id: @key_id)
      open_auth_file('r+') do |f|
        while line = f.gets # rubocop:disable Lint/AssignmentInCondition
          next unless line.start_with?("command=\"#{self.class.command_key(@key_id)}\"")
          f.seek(-line.length, IO::SEEK_CUR)
          # Overwrite the line with #'s. Because the 'line' variable contains
          # a terminating '\n', we write line.length - 1 '#' characters.
          f.write('#' * (line.length - 1))
        end
{'fraction_non_alphanumeric': 0.13067524115755627, 'fraction_numerical': 0.005916398713826367, 'mean_word_length': 2.490125673249551, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4966690', 'n_tokens_mistral': 2733, 'n_tokens_neox': 2658, 'n_words': 580}