<commit_before>require 'chef/knife'
require 'chef/knife/ec_base'

class Chef
  class Knife
    class EcRestore < Chef::Knife

      include Knife::EcBase

      banner "knife ec restore DIRECTORY"

      option :overwrite_pivotal,
        :long => '--overwrite-pivotal',
        :boolean => true,
        :default => false,
        :description => "Whether to overwrite pivotal's key.  Once this is done, future requests will fail until you fix the private key."

      option :skip_users,
        :long => "--skip-users",
        :description => "Skip restoring users"

      deps do
        require 'chef/json_compat'
        require 'chef/chef_fs/config'
        require 'chef/chef_fs/file_system'
        require 'chef/chef_fs/file_pattern'
        # Work around bug in chef_fs
        require 'chef/chef_fs/command_line'
        require 'chef/chef_fs/data_handler/acl_data_handler'
        require 'securerandom'
        require 'chef/chef_fs/parallelizer'
        require 'chef/tsorter'
        require 'chef/server'
      end

      def run
        set_dest_dir_from_args!
        set_client_config!
        ensure_webui_key_exists!
        set_skip_user_acl!

        warn_on_incorrect_clients_group(dest_dir, "restore")

        restore_users unless config[:skip_users]
        restore_user_sql if config[:with_user_sql]

        for_each_organization do |orgname|
          ui.msg "Restoring organization[#{orgname}]"
          create_organization(orgname)
          restore_open_invitations(orgname)
          add_users_to_org(orgname)
          upload_org_data(orgname)
        end

        restore_key_sql if config[:with_key_sql]

        if config[:skip_useracl]
          ui.warn("Skipping user ACL update. To update user ACLs, remove --skip-useracl or upgrade your Enterprise Chef Server.")
        else
          restore_user_acls
        end

        completion_banner
      end

      def create_organization(orgname)
        org = JSONCompat.from_json(File.read("#{dest_dir}/organizations/#{orgname}/org.json"))
        rest.post('organizations', org)
      rescue Net::HTTPServerException => ex
        if ex.response.code == "409"
          rest.put("organizations/#{orgname}", org)
        else
          knife_ec_error_handler.add(ex)
        end
      end

      def restore_open_invitations(orgname)
        invitations = JSONCompat.from_json(File.read("#{dest_dir}/organizations/#{orgname}/invitations.json"))
        invitations.each do |invitation|
          begin
            rest.post("organizations/#{orgname}/association_requests", { 'user' => invitation['username'] })
          rescue Net::HTTPServerException => ex
            if ex.response.code != "409"
              ui.error("Cannot create invitation #{invitation['id']}")
              knife_ec_error_handler.add(ex)
            end
          end
        end
      end

      def add_users_to_org(orgname)
        members = JSONCompat.from_json(File.read("#{dest_dir}/organizations/#{orgname}/members.json"))
        members.each do |member|
          username = member['user']['username']
          begin
            response = rest.post("organizations/#{orgname}/association_requests", { 'user' => username })
            association_id = response["uri"].split("/").last
            rest.put("users/#{username}/association_requests/#{association_id}", { 'response' => 'accept' })
          rescue Net::HTTPServerException => ex
            knife_ec_error_handler.add(ex) if ex.response.code != "409"
          end
        end
      end

      def restore_user_acls
        ui.msg "Restoring user ACLs"
        for_each_user do |name|
          user_acl = JSONCompat.from_json(File.read("#{dest_dir}/user_acls/#{name}.json"))
          put_acl(user_acl_rest, "users/#{name}/_acl", user_acl)
        end
      end

      def for_each_user
        Dir.foreach("#{dest_dir}/users") do |filename|
          next if filename !~ /(.+)\.json/
          name = $1
          if name == 'pivotal' && !config[:overwrite_pivotal]
            ui.warn("Skipping pivotal user.  To overwrite pivotal, pass --overwrite-pivotal.")
            next
          end
          yield name
        end
      end

      def for_each_organization
        Dir.foreach("#{dest_dir}/organizations") do |name|
          next if name == '..' || name == '.' || !File.directory?("#{dest_dir}/organizations/#{name}")
          next unless (config[:org].nil? || config[:org] == name)
          yield name
        end
      end

      def restore_users
        ui.msg "Restoring users"
        for_each_user do |name|
          user = JSONCompat.from_json(File.read("#{dest_dir}/users/#{name}.json"))
          begin
            # Supply password for new user
            user_with_password = <PASSWORD>
            user_with_password['password'] = <PASSWORD>
            rest.post('users', user_with_password)
          rescue Net::HTTPServerException => ex
            if ex.response.code == "409"
              rest.put("users/#{name}", user)
              next
            end
            knife_ec_error_handler.add(ex)
          end
        end
        purge_users_on_restore
      end

      def users_for_purge
        purge_list = remote_user_list - local_user_list
        # failsafe - don't delete pivotal
        purge_list -= ['pivotal']
        purge_list.each do |user|
          yield user
        end
      end

      def purge_users_on_restore
        return unless config[:purge]
        users_for_purge do |user|
          ui.msg "Deleting user #{user} from remote (purge is on)"
          begin
            rest.delete("/users/#{user}")
          rescue Net::HTTPServerException => e
            ui.warn "Failed deleting user #{user} from remote #{e}"
          end
        end
      end

      def ec_key_import
        @ec_key_import ||= begin
                             require 'chef/knife/ec_key_import'
                             k = Chef::Knife::EcKeyImport.new
                             k.name_args = ["#{dest_dir}/key_dump.json", "#{dest_dir}/key_table_dump.json"]
                             k.config[:skip_pivotal] = true
                             k.config[:skip_ids] = true
                             k.config[:sql_host] = config[:sql_host]
                             k.config[:sql_port] = config[:sql_port]
                             k.config[:sql_user] = config[:sql_user]
                             k.config[:sql_password] = config[:sql_password]
                             k
                           end
      end

      def restore_user_sql
        k = ec_key_import<commit_msg>Leverage config[:skip_ids] for key import instead of hardcoding the value in ec_restore.rb<commit_after>require 'chef/knife'
require 'chef/knife/ec_base'

class Chef
  class Knife
    class EcRestore < Chef::Knife

      include Knife::EcBase

      banner "knife ec restore DIRECTORY"

      option :overwrite_pivotal,
        :long => '--overwrite-pivotal',
        :boolean => true,
        :default => false,
        :description => "Whether to overwrite pivotal's key.  Once this is done, future requests will fail until you fix the private key."

      option :skip_users,
        :long => "--skip-users",
        :description => "Skip restoring users"

      option :skip_ids,
        :long => "--[no-]skip-user-ids",
        :default => true,
        :boolean => true,
        :description => "Upload user ids."

      deps do
        require 'chef/json_compat'
        require 'chef/chef_fs/config'
        require 'chef/chef_fs/file_system'
        require 'chef/chef_fs/file_pattern'
        # Work around bug in chef_fs
        require 'chef/chef_fs/command_line'
        require 'chef/chef_fs/data_handler/acl_data_handler'
        require 'securerandom'
        require 'chef/chef_fs/parallelizer'
        require 'chef/tsorter'
        require 'chef/server'
      end

      def run
        set_dest_dir_from_args!
        set_client_config!
        ensure_webui_key_exists!
        set_skip_user_acl!

        warn_on_incorrect_clients_group(dest_dir, "restore")

        restore_users unless config[:skip_users]
        restore_user_sql if config[:with_user_sql]

        for_each_organization do |orgname|
          ui.msg "Restoring organization[#{orgname}]"
          create_organization(orgname)
          restore_open_invitations(orgname)
          add_users_to_org(orgname)
          upload_org_data(orgname)
        end

        restore_key_sql if config[:with_key_sql]

        if config[:skip_useracl]
          ui.warn("Skipping user ACL update. To update user ACLs, remove --skip-useracl or upgrade your Enterprise Chef Server.")
        else
          restore_user_acls
        end

        completion_banner
      end

      def create_organization(orgname)
        org = JSONCompat.from_json(File.read("#{dest_dir}/organizations/#{orgname}/org.json"))
        rest.post('organizations', org)
      rescue Net::HTTPServerException => ex
        if ex.response.code == "409"
          rest.put("organizations/#{orgname}", org)
        else
          knife_ec_error_handler.add(ex)
        end
      end

      def restore_open_invitations(orgname)
        invitations = JSONCompat.from_json(File.read("#{dest_dir}/organizations/#{orgname}/invitations.json"))
        invitations.each do |invitation|
          begin
            rest.post("organizations/#{orgname}/association_requests", { 'user' => invitation['username'] })
          rescue Net::HTTPServerException => ex
            if ex.response.code != "409"
              ui.error("Cannot create invitation #{invitation['id']}")
              knife_ec_error_handler.add(ex)
            end
          end
        end
      end

      def add_users_to_org(orgname)
        members = JSONCompat.from_json(File.read("#{dest_dir}/organizations/#{orgname}/members.json"))
        members.each do |member|
          username = member['user']['username']
          begin
            response = rest.post("organizations/#{orgname}/association_requests", { 'user' => username })
            association_id = response["uri"].split("/").last
            rest.put("users/#{username}/association_requests/#{association_id}", { 'response' => 'accept' })
          rescue Net::HTTPServerException => ex
            knife_ec_error_handler.add(ex) if ex.response.code != "409"
          end
        end
      end

      def restore_user_acls
        ui.msg "Restoring user ACLs"
        for_each_user do |name|
          user_acl = JSONCompat.from_json(File.read("#{dest_dir}/user_acls/#{name}.json"))
          put_acl(user_acl_rest, "users/#{name}/_acl", user_acl)
        end
      end

      def for_each_user
        Dir.foreach("#{dest_dir}/users") do |filename|
          next if filename !~ /(.+)\.json/
          name = $1
          if name == 'pivotal' && !config[:overwrite_pivotal]
            ui.warn("Skipping pivotal user.  To overwrite pivotal, pass --overwrite-pivotal.")
            next
          end
          yield name
        end
      end

      def for_each_organization
        Dir.foreach("#{dest_dir}/organizations") do |name|
          next if name == '..' || name == '.' || !File.directory?("#{dest_dir}/organizations/#{name}")
          next unless (config[:org].nil? || config[:org] == name)
          yield name
        end
      end

      def restore_users
        ui.msg "Restoring users"
        for_each_user do |name|
          user = JSONCompat.from_json(File.read("#{dest_dir}/users/#{name}.json"))
          begin
            # Supply password for new user
            user_with_password = user.dup
            user_with_password['password'] = <PASSWORD>
            rest.post('users', user_with_password)
          rescue Net::HTTPServerException => ex
            if ex.response.code == "409"
              rest.put("users/#{name}", user)
              next
            end
            knife_ec_error_handler.add(ex)
          end
        end
        purge_users_on_restore
      end

      def users_for_purge
        purge_list = remote_user_list - local_user_list
        # failsafe - don't delete pivotal
        purge_list -= ['pivotal']
        purge_list.each do |user|
          yield user
        end
      end

      def purge_users_on_restore
        return unless config[:purge]
        users_for_purge do |user|
          ui.msg "Deleting user #{user} from remote (purge is on)"
          begin
            rest.delete("/users/#{user}")
          rescue Net::HTTPServerException => e
            ui.warn "Failed deleting user #{user} from remote #{e}"
          end
        end
      end

      def ec_key_import
        @ec_key_import ||= begin
                             require 'chef/knife/ec_key_import'
                             k = Chef::Knife::EcKeyImport.new
                             k.name_args = ["#{dest_dir}/key_dump.json", "#{dest_dir}/key_table_dump.json"]
                             k.config[:skip_pivotal] = true
                             k.config[:skip_ids] = config[:skip_ids]
                             k.config[:sql_host] = config[:sql_host]
                             k.config[:sql_port] = config[:sql_port]
                             k.config[:sql_user] = config[:sql_user]
                             k.config[:sql_password] = config[:sql_password]
                             k
                           end
      end

      def restore_user_sql
        k = ec_key_import
{'fraction_non_alphanumeric': 0.10517305545217714, 'fraction_numerical': 0.001935243766282099, 'mean_word_length': 1.9355473017260214, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 38, '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': '2240126', 'n_tokens_mistral': 3958, 'n_tokens_neox': 3865, 'n_words': 798}