opennebula 3.2.1-2 / usr / bin / onequota

#!/usr/bin/env ruby

# -------------------------------------------------------------------------- #
# Copyright 2002-2012, OpenNebula Project Leads (OpenNebula.org)             #
#                                                                            #
# Licensed under the Apache License, Version 2.0 (the "License"); you may    #
# not use this file except in compliance with the License. You may obtain    #
# a copy of the License at                                                   #
#                                                                            #
# http://www.apache.org/licenses/LICENSE-2.0                                 #
#                                                                            #
# Unless required by applicable law or agreed to in writing, software        #
# distributed under the License is distributed on an "AS IS" BASIS,          #
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   #
# See the License for the specific language governing permissions and        #
# limitations under the License.                                             #
#--------------------------------------------------------------------------- #

ONE_LOCATION=ENV["ONE_LOCATION"]

if !ONE_LOCATION
    RUBY_LIB_LOCATION="/usr/lib/one/ruby"
else
    RUBY_LIB_LOCATION=ONE_LOCATION+"/lib/ruby"
end

$: << RUBY_LIB_LOCATION
$: << RUBY_LIB_LOCATION+"/cli"

require 'command_parser'
require 'cli_helper'
require 'one_helper'
require 'quota'

require 'pp'

QUOTA_KEYS = Quota::DB_QUOTA_SCHEMA.keys

def show_table(values, usage=nil)
    size = usage ? 12 : 8

    values.sort!{|v1, v2| v1[:UID]<=>v2[:UID] }

    table=CLIHelper::ShowTable.new(nil, self) do
        column :UID, "ONE identifier for the User", :size=>4 do |d|
            d[:UID]
        end

        QUOTA_KEYS.each { |key|
            column key, "#{key} quota", :size=>size do |d|
                if usage
                    "#{usage[key].to_i}/#{d[key].to_i}"
                else
                    "#{d[key].to_i}"
                end
            end
        }

        default :UID, *QUOTA_KEYS
    end

    table.show(values)
end

cmd=CommandParser::CmdParser.new(ARGV) do
    usage "`onequota` <command> [<args>] [<options>]"
    version OpenNebulaHelper::ONE_VERSION

    quota = Quota.new

    ########################################################################
    # Global Options
    ########################################################################
    set :option, CommandParser::OPTIONS

    ########################################################################
    # Argument Formats
    ########################################################################
    quota_list_desc = <<-EOT.unindent
        List of quota keys, comma separated.
        Available quotas: #{QUOTA_KEYS.join(', ')}
    EOT

    set :format, :quota_list, quota_list_desc do |arg|
        arg_list = arg.split(',')

        rc = nil
        arg_list.collect! do |elem|
          sym = elem.upcase.to_sym
          if !QUOTA_KEYS.include?(sym)
            rc = -1, "#{elem} is not a valid quota"
            break
          end
          sym
        end

        rc ? rc : [0, arg_list]
    end

    set :format, :value_list, "List of quota values, comma separated." do |arg|
        arg_list = arg.split(',')
        arg_list.map! {|a| a.to_i }
        [0, arg_list]
    end

    set :format, :userid, OpenNebulaHelper.rname_to_id_desc("USER") do |arg|
        OpenNebulaHelper.rname_to_id(arg, "USER")
    end

    ########################################################################
    # Commands
    ########################################################################
    set_desc = <<-EOT.unindent
        Set a quota for a given user.
        Examples:
          onequota set 3 cpu 12
          onequota set 4 cpu,memory,storage 8,4096,10000
    EOT

    command :set, set_desc, :userid, :quota_list, :value_list do
        user_id, keys, values = args

        if keys.length != values.length
            exit_with_code -1, "The number of keys and values does not match"
        end

        values_hash = Hash.new
        keys.each_with_index { |k,i|
            values_hash[k] = values[i]
        }

        quota.set_quota(user_id, values_hash)
        exit_with_code 0
    end

    ########################################################################
    unset_desc = <<-EOT.unindent
        Unset a quota for a given user.
        Examples:
          onequota unset 3 cpu
          onequota unset 4 cpu,memory,storage
    EOT

    command :unset, unset_desc, :userid, :quota_list do
        user_id, keys = args

        values_hash = Hash.new
        keys.each_with_index { |k,i|
            values_hash[k] = 0
        }

        quota.set_quota(user_id, values_hash)
        exit_with_code 0
    end

    ########################################################################
    delete_desc = "Delete the defined quotas for the given user"

    command :delete, delete_desc, :userid do
        quota.delete_quota(args[0])
        exit_with_code 0
    end

    ########################################################################
    show_desc = "Show the user's quota and usage. (usage/quota)"

    FORCE={
        :name  => "force",
        :short => "-f",
        :large => "--force",
        :description => "Force the usage calculation instead of using the cache"
    }

    command :show, show_desc, :userid, :options=>[FORCE] do
        user_usage = quota.get_usage(args[0],nil,options[:force])
        user_quota = quota.get_quota(args[0])

        show_table([user_quota], user_usage)
        exit_with_code 0
    end

    ########################################################################
    list_desc = "List the defined quotas for all the users"

    command :list, list_desc do
        show_table(quota.get_quota)
        exit_with_code 0
    end
end