xref: /trafficserver/src/tscore/ArgParser.cc (revision 4a43f53e)
1 /** @file
2 
3   Powerful and easy-to-use command line parsing for ATS
4 
5   @section license License
6 
7   Licensed to the Apache Software Foundation (ASF) under one
8   or more contributor license agreements.  See the NOTICE file
9   distributed with this work for additional information
10   regarding copyright ownership.  The ASF licenses this file
11   to you under the Apache License, Version 2.0 (the
12   "License"); you may not use this file except in compliance
13   with the License.  You may obtain a copy of the License at
14 
15       http://www.apache.org/licenses/LICENSE-2.0
16 
17   Unless required by applicable law or agreed to in writing, software
18   distributed under the License is distributed on an "AS IS" BASIS,
19   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20   See the License for the specific language governing permissions and
21   limitations under the License.
22  */
23 
24 #include "tscore/ArgParser.h"
25 #include "tscore/ink_file.h"
26 #include "tscore/I_Version.h"
27 
28 #include <iostream>
29 #include <set>
30 #include <sstream>
31 #include <utility>
32 #include <sysexits.h>
33 
34 std::string global_usage;
35 std::string parser_program_name;
36 std::string default_command;
37 
38 // by default return EX_USAGE(64) when usage is called.
39 // if -h or --help is called specifically, return 0
40 int usage_return_code = EX_USAGE;
41 
42 namespace ts
43 {
ArgParser()44 ArgParser::ArgParser() {}
45 
ArgParser(std::string const & name,std::string const & description,std::string const & envvar,unsigned arg_num,Function const & f)46 ArgParser::ArgParser(std::string const &name, std::string const &description, std::string const &envvar, unsigned arg_num,
47                      Function const &f)
48 {
49   // initialize _top_level_command according to the provided message
50   _top_level_command = ArgParser::Command(name, description, envvar, arg_num, f);
51 }
52 
~ArgParser()53 ArgParser::~ArgParser() {}
54 
55 // add new options with args
56 ArgParser::Command &
add_option(std::string const & long_option,std::string const & short_option,std::string const & description,std::string const & envvar,unsigned arg_num,std::string const & default_value,std::string const & key)57 ArgParser::add_option(std::string const &long_option, std::string const &short_option, std::string const &description,
58                       std::string const &envvar, unsigned arg_num, std::string const &default_value, std::string const &key)
59 {
60   return _top_level_command.add_option(long_option, short_option, description, envvar, arg_num, default_value, key);
61 }
62 
63 // add sub-command with only function
64 ArgParser::Command &
add_command(std::string const & cmd_name,std::string const & cmd_description,Function const & f,std::string const & key)65 ArgParser::add_command(std::string const &cmd_name, std::string const &cmd_description, Function const &f, std::string const &key)
66 {
67   return _top_level_command.add_command(cmd_name, cmd_description, f, key);
68 }
69 
70 // add sub-command without args and function
71 ArgParser::Command &
add_command(std::string const & cmd_name,std::string const & cmd_description,std::string const & cmd_envvar,unsigned cmd_arg_num,Function const & f,std::string const & key)72 ArgParser::add_command(std::string const &cmd_name, std::string const &cmd_description, std::string const &cmd_envvar,
73                        unsigned cmd_arg_num, Function const &f, std::string const &key)
74 {
75   return _top_level_command.add_command(cmd_name, cmd_description, cmd_envvar, cmd_arg_num, f, key);
76 }
77 
78 void
add_global_usage(std::string const & usage)79 ArgParser::add_global_usage(std::string const &usage)
80 {
81   global_usage = usage;
82 }
83 
84 void
help_message(std::string_view err) const85 ArgParser::help_message(std::string_view err) const
86 {
87   _top_level_command.help_message(err);
88 }
89 
90 // a graceful way to output help message
91 void
help_message(std::string_view err) const92 ArgParser::Command::help_message(std::string_view err) const
93 {
94   if (!err.empty()) {
95     std::cout << "Error: " << err << std::endl;
96   }
97   // output global usage
98   if (global_usage.size() > 0) {
99     std::cout << "\nUsage: " + global_usage << std::endl;
100   }
101   // output subcommands
102   std::cout << "\nCommands ---------------------- Description -----------------------" << std::endl;
103   std::string prefix = "";
104   output_command(std::cout, prefix);
105   // output options
106   if (_option_list.size() > 0) {
107     std::cout << "\nOptions ======================= Default ===== Description =============" << std::endl;
108     output_option();
109   }
110   // output example usage
111   if (!_example_usage.empty()) {
112     std::cout << "\nExample Usage: " << _example_usage << std::endl;
113   }
114   // standard return code
115   exit(usage_return_code);
116 }
117 
118 void
version_message() const119 ArgParser::Command::version_message() const
120 {
121   // unified version message of ATS
122   AppVersionInfo appVersionInfo;
123   appVersionInfo.setup(PACKAGE_NAME, _name.c_str(), PACKAGE_VERSION, __DATE__, __TIME__, BUILD_MACHINE, BUILD_PERSON, "");
124   ink_fputln(stdout, appVersionInfo.FullVersionInfoStr);
125   exit(0);
126 }
127 
128 void
set_default_command(std::string const & cmd)129 ArgParser::set_default_command(std::string const &cmd)
130 {
131   if (default_command.empty()) {
132     if (_top_level_command._subcommand_list.find(cmd) == _top_level_command._subcommand_list.end()) {
133       std::cerr << "Error: Default command " << cmd << "not found" << std::endl;
134       exit(1);
135     }
136     default_command = cmd;
137   } else if (cmd != default_command) {
138     std::cerr << "Error: Default command " << default_command << "already existed" << std::endl;
139     exit(1);
140   }
141 }
142 
143 // Top level call of parsing
144 Arguments
parse(const char ** argv)145 ArgParser::parse(const char **argv)
146 {
147   // deal with argv first
148   int size = 0;
149   _argv.clear();
150   while (argv[size]) {
151     _argv.push_back(argv[size]);
152     size++;
153   }
154   if (size == 0) {
155     std::cout << "Error: invalid argv provided" << std::endl;
156     exit(1);
157   }
158   // the name of the program only
159   _argv[0]                 = _argv[0].substr(_argv[0].find_last_of('/') + 1);
160   _top_level_command._name = _argv[0];
161   _top_level_command._key  = _argv[0];
162   parser_program_name      = _argv[0];
163   Arguments ret; // the parsed arg object to return
164   AP_StrVec args = _argv;
165   // call the recursive parse method in Command
166   if (!_top_level_command.parse(ret, args)) {
167     // deal with default command
168     if (!default_command.empty()) {
169       args = _argv;
170       args.insert(args.begin() + 1, default_command);
171       _top_level_command.parse(ret, args);
172     }
173   };
174   // if there is anything left, then output usage
175   if (!args.empty()) {
176     std::string msg = "Unknown command, option or args:";
177     for (const auto &it : args) {
178       msg = msg + " '" + it + "'";
179     }
180     // find the correct level to output help message
181     ArgParser::Command *command = &_top_level_command;
182     for (unsigned i = 1; i < _argv.size(); i++) {
183       auto it = command->_subcommand_list.find(_argv[i]);
184       if (it == command->_subcommand_list.end()) {
185         break;
186       }
187       command = &it->second;
188     }
189     command->help_message(msg);
190   }
191   return ret;
192 }
193 
194 ArgParser::Command &
require_commands()195 ArgParser::require_commands()
196 {
197   return _top_level_command.require_commands();
198 }
199 
200 void
set_error(std::string e)201 ArgParser::set_error(std::string e)
202 {
203   _error_msg = std::move(e);
204 }
205 
206 std::string
get_error() const207 ArgParser::get_error() const
208 {
209   return _error_msg;
210 }
211 
212 //=========================== Command class ================================
Command()213 ArgParser::Command::Command() {}
214 
~Command()215 ArgParser::Command::~Command() {}
216 
Command(std::string const & name,std::string const & description,std::string const & envvar,unsigned arg_num,Function const & f,std::string const & key)217 ArgParser::Command::Command(std::string const &name, std::string const &description, std::string const &envvar, unsigned arg_num,
218                             Function const &f, std::string const &key)
219   : _name(name), _description(description), _arg_num(arg_num), _envvar(envvar), _f(f), _key(key)
220 {
221 }
222 
223 // check if this is a valid option before adding
224 void
check_option(std::string const & long_option,std::string const & short_option,std::string const & key) const225 ArgParser::Command::check_option(std::string const &long_option, std::string const &short_option, std::string const &key) const
226 {
227   if (long_option.size() < 3 || long_option[0] != '-' || long_option[1] != '-') {
228     // invalid name
229     std::cerr << "Error: invalid long option added: '" + long_option + "'" << std::endl;
230     exit(1);
231   }
232   if (short_option.size() > 2 || (short_option.size() > 0 && short_option[0] != '-')) {
233     // invalid short option
234     std::cerr << "Error: invalid short option added: '" + short_option + "'" << std::endl;
235     exit(1);
236   }
237   // find if existing in option list
238   if (_option_list.find(long_option) != _option_list.end()) {
239     std::cerr << "Error: long option '" + long_option + "' already existed" << std::endl;
240     exit(1);
241   } else if (_option_map.find(short_option) != _option_map.end()) {
242     std::cerr << "Error: short option '" + short_option + "' already existed" << std::endl;
243     exit(1);
244   }
245 }
246 
247 // check if this is a valid command before adding
248 void
check_command(std::string const & name,std::string const & key) const249 ArgParser::Command::check_command(std::string const &name, std::string const &key) const
250 {
251   if (name.empty()) {
252     // invalid name
253     std::cerr << "Error: empty command cannot be added" << std::endl;
254     exit(1);
255   }
256   // find if existing in subcommand list
257   if (_subcommand_list.find(name) != _subcommand_list.end()) {
258     std::cerr << "Error: command already exists: '" + name + "'" << std::endl;
259     exit(1);
260   }
261 }
262 
263 // add new options with args
264 ArgParser::Command &
add_option(std::string const & long_option,std::string const & short_option,std::string const & description,std::string const & envvar,unsigned arg_num,std::string const & default_value,std::string const & key)265 ArgParser::Command::add_option(std::string const &long_option, std::string const &short_option, std::string const &description,
266                                std::string const &envvar, unsigned arg_num, std::string const &default_value,
267                                std::string const &key)
268 {
269   std::string lookup_key = key.empty() ? long_option.substr(2) : key;
270   check_option(long_option, short_option, lookup_key);
271   _option_list[long_option] = {long_option, short_option == "-" ? "" : short_option, description, envvar, arg_num, default_value,
272                                lookup_key};
273   if (short_option != "-" && !short_option.empty()) {
274     _option_map[short_option] = long_option;
275   }
276   return *this;
277 }
278 
279 // add sub-command with only function
280 ArgParser::Command &
add_command(std::string const & cmd_name,std::string const & cmd_description,Function const & f,std::string const & key)281 ArgParser::Command::add_command(std::string const &cmd_name, std::string const &cmd_description, Function const &f,
282                                 std::string const &key)
283 {
284   std::string lookup_key = key.empty() ? cmd_name : key;
285   check_command(cmd_name, lookup_key);
286   _subcommand_list[cmd_name] = ArgParser::Command(cmd_name, cmd_description, "", 0, f, lookup_key);
287   return _subcommand_list[cmd_name];
288 }
289 
290 // add sub-command without args and function
291 ArgParser::Command &
add_command(std::string const & cmd_name,std::string const & cmd_description,std::string const & cmd_envvar,unsigned cmd_arg_num,Function const & f,std::string const & key)292 ArgParser::Command::add_command(std::string const &cmd_name, std::string const &cmd_description, std::string const &cmd_envvar,
293                                 unsigned cmd_arg_num, Function const &f, std::string const &key)
294 {
295   std::string lookup_key = key.empty() ? cmd_name : key;
296   check_command(cmd_name, lookup_key);
297   _subcommand_list[cmd_name] = ArgParser::Command(cmd_name, cmd_description, cmd_envvar, cmd_arg_num, f, lookup_key);
298   return _subcommand_list[cmd_name];
299 }
300 
301 ArgParser::Command &
add_example_usage(std::string const & usage)302 ArgParser::Command::add_example_usage(std::string const &usage)
303 {
304   _example_usage = usage;
305   return *this;
306 }
307 
308 // method used by help_message()
309 void
output_command(std::ostream & out,std::string const & prefix) const310 ArgParser::Command::output_command(std::ostream &out, std::string const &prefix) const
311 {
312   if (_name != parser_program_name) {
313     // a nicely formatted way to output command usage
314     std::string msg = prefix + _name;
315     // nicely formatted output
316     if (!_description.empty()) {
317       if (INDENT_ONE - static_cast<int>(msg.size()) < 0) {
318         // if the command msg is too long
319         std::cout << msg << "\n" << std::string(INDENT_ONE, ' ') << _description << std::endl;
320       } else {
321         std::cout << msg << std::string(INDENT_ONE - msg.size(), ' ') << _description << std::endl;
322       }
323     }
324   }
325   // recursive call
326   for (const auto &it : _subcommand_list) {
327     it.second.output_command(out, "  " + prefix);
328   }
329 }
330 
331 // a nicely formatted way to output option message for help.
332 void
output_option() const333 ArgParser::Command::output_option() const
334 {
335   for (const auto &it : _option_list) {
336     std::string msg;
337     if (!it.second.short_option.empty()) {
338       msg = it.second.short_option + ", ";
339     }
340     msg += it.first;
341     unsigned num = it.second.arg_num;
342     if (num != 0) {
343       if (num == 1) {
344         msg = msg + " <arg>";
345       } else if (num == MORE_THAN_ZERO_ARG_N) {
346         msg = msg + " [<arg> ...]";
347       } else if (num == MORE_THAN_ONE_ARG_N) {
348         msg = msg + " <arg> ...";
349       } else {
350         msg = msg + " <arg1> ... <arg" + std::to_string(num) + ">";
351       }
352     }
353     if (!it.second.default_value.empty()) {
354       if (INDENT_ONE - static_cast<int>(msg.size()) < 0) {
355         msg = msg + "\n" + std::string(INDENT_ONE, ' ') + it.second.default_value;
356       } else {
357         msg = msg + std::string(INDENT_ONE - msg.size(), ' ') + it.second.default_value;
358       }
359     }
360     if (!it.second.description.empty()) {
361       if (INDENT_TWO - static_cast<int>(msg.size()) < 0) {
362         std::cout << msg << "\n" << std::string(INDENT_TWO, ' ') << it.second.description << std::endl;
363       } else {
364         std::cout << msg << std::string(INDENT_TWO - msg.size(), ' ') << it.second.description << std::endl;
365       }
366     }
367   }
368 }
369 
370 // helper method to handle the arguments and put them nicely in arguments
371 // can be switched to ts::errata
372 static std::string
handle_args(Arguments & ret,AP_StrVec & args,std::string const & name,unsigned arg_num,unsigned & index)373 handle_args(Arguments &ret, AP_StrVec &args, std::string const &name, unsigned arg_num, unsigned &index)
374 {
375   ArgumentData data;
376   ret.append(name, data);
377   // handle the args
378   if (arg_num == MORE_THAN_ZERO_ARG_N || arg_num == MORE_THAN_ONE_ARG_N) {
379     // infinite arguments
380     if (arg_num == MORE_THAN_ONE_ARG_N && args.size() <= index + 1) {
381       return "at least one argument expected by " + name;
382     }
383     for (unsigned j = index + 1; j < args.size(); j++) {
384       ret.append_arg(name, args[j]);
385     }
386     args.erase(args.begin() + index, args.end());
387     return "";
388   }
389   // finite number of argument handling
390   for (unsigned j = 0; j < arg_num; j++) {
391     if (args.size() < index + j + 2 || args[index + j + 1].empty()) {
392       return std::to_string(arg_num) + " argument(s) expected by " + name;
393     }
394     ret.append_arg(name, args[index + j + 1]);
395   }
396   // erase the used arguments and append the data to the return structure
397   args.erase(args.begin() + index, args.begin() + index + arg_num + 1);
398   index -= 1;
399   return "";
400 }
401 
402 // Append the args of option to parsed data. Return true if there is any option called
403 void
append_option_data(Arguments & ret,AP_StrVec & args,int index)404 ArgParser::Command::append_option_data(Arguments &ret, AP_StrVec &args, int index)
405 {
406   std::map<std::string, unsigned> check_map;
407   for (unsigned i = index; i < args.size(); i++) {
408     // find matches of the arg
409     if (args[i][0] == '-' && args[i][1] == '-' && args[i].find('=') != std::string::npos) {
410       // deal with --args=
411       std::string option_name = args[i].substr(0, args[i].find_first_of('='));
412       std::string value       = args[i].substr(args[i].find_last_of('=') + 1);
413       if (value.empty()) {
414         help_message("missing argument for '" + option_name + "'");
415       }
416       auto it = _option_list.find(option_name);
417       if (it != _option_list.end()) {
418         ArgParser::Option cur_option = it->second;
419         // handle environment variable
420         if (!cur_option.envvar.empty()) {
421           ret.set_env(cur_option.key, getenv(cur_option.envvar.c_str()) ? getenv(cur_option.envvar.c_str()) : "");
422         }
423         ret.append_arg(cur_option.key, value);
424         check_map[cur_option.long_option] += 1;
425         args.erase(args.begin() + i);
426         i -= 1;
427       }
428     } else {
429       // output version message
430       if ((args[i] == "--version" || args[i] == "-V") && _option_list.find("--version") != _option_list.end()) {
431         version_message();
432       }
433       // output help message
434       if ((args[i] == "--help" || args[i] == "-h") && _option_list.find("--help") != _option_list.end()) {
435         ArgParser::Command *command = this;
436         // find the correct level to output help message
437         for (unsigned i = 1; i < args.size(); i++) {
438           auto it = command->_subcommand_list.find(args[i]);
439           if (it == command->_subcommand_list.end()) {
440             break;
441           }
442           command = &it->second;
443         }
444         usage_return_code = 0;
445         command->help_message();
446       }
447       // deal with normal --arg val1 val2 ...
448       auto long_it  = _option_list.find(args[i]);
449       auto short_it = _option_map.find(args[i]);
450       // long option match or short option match
451       if (long_it != _option_list.end() || short_it != _option_map.end()) {
452         ArgParser::Option cur_option;
453         if (long_it != _option_list.end()) {
454           cur_option = long_it->second;
455         } else {
456           cur_option = _option_list.at(short_it->second);
457         }
458         // handle the arguments
459         std::string err = handle_args(ret, args, cur_option.key, cur_option.arg_num, i);
460         if (!err.empty()) {
461           help_message(err);
462         }
463         // handle environment variable
464         if (!cur_option.envvar.empty()) {
465           ret.set_env(cur_option.key, getenv(cur_option.envvar.c_str()) ? getenv(cur_option.envvar.c_str()) : "");
466         }
467       }
468     }
469   }
470   // check for wrong number of arguments for --arg=...
471   for (const auto &it : check_map) {
472     unsigned num = _option_list.at(it.first).arg_num;
473     if (num != it.second && num < MORE_THAN_ONE_ARG_N) {
474       help_message(std::to_string(_option_list.at(it.first).arg_num) + " arguments expected by " + it.first);
475     }
476   }
477   // put in the default value of options
478   for (const auto &it : _option_list) {
479     if (!it.second.default_value.empty() && ret.get(it.second.key).empty()) {
480       std::istringstream ss(it.second.default_value);
481       std::string token;
482       while (std::getline(ss, token, ' ')) {
483         ret.append_arg(it.second.key, token);
484       }
485     }
486   }
487 }
488 
489 // Main recursive logic of Parsing
490 bool
parse(Arguments & ret,AP_StrVec & args)491 ArgParser::Command::parse(Arguments &ret, AP_StrVec &args)
492 {
493   bool command_called = false;
494   // iterate through all arguments
495   for (unsigned i = 0; i < args.size(); i++) {
496     if (_name == args[i]) {
497       command_called = true;
498       // handle the option
499       append_option_data(ret, args, i);
500       // handle the action
501       if (_f) {
502         ret._action = _f;
503       }
504       std::string err = handle_args(ret, args, _key, _arg_num, i);
505       if (!err.empty()) {
506         help_message(err);
507       }
508       // set ENV var
509       if (!_envvar.empty()) {
510         ret.set_env(_key, getenv(_envvar.c_str()) ? getenv(_envvar.c_str()) : "");
511       }
512       break;
513     }
514   }
515   if (command_called) {
516     bool flag = false;
517     // recursively call subcommand
518     for (auto &it : _subcommand_list) {
519       if (it.second.parse(ret, args)) {
520         flag = true;
521         break;
522       }
523     }
524     // check for command required
525     if (!flag && _command_required) {
526       help_message("No subcommand found for " + _name);
527     }
528     if (_name == parser_program_name) {
529       // if we are at the top level
530       return flag;
531     }
532   }
533   return command_called;
534 }
535 
536 ArgParser::Command &
require_commands()537 ArgParser::Command::require_commands()
538 {
539   _command_required = true;
540   return *this;
541 }
542 
543 ArgParser::Command &
set_default()544 ArgParser::Command::set_default()
545 {
546   default_command = _name;
547   return *this;
548 }
549 
550 //=========================== Arguments class ================================
551 
Arguments()552 Arguments::Arguments() {}
~Arguments()553 Arguments::~Arguments() {}
554 
555 ArgumentData
get(std::string const & name)556 Arguments::get(std::string const &name)
557 {
558   if (_data_map.find(name) != _data_map.end()) {
559     _data_map[name]._is_called = true;
560     return _data_map[name];
561   }
562   return ArgumentData();
563 }
564 
565 void
append(std::string const & key,ArgumentData const & value)566 Arguments::append(std::string const &key, ArgumentData const &value)
567 {
568   // perform overwrite for now
569   _data_map[key] = value;
570 }
571 
572 void
append_arg(std::string const & key,std::string const & value)573 Arguments::append_arg(std::string const &key, std::string const &value)
574 {
575   _data_map[key]._values.push_back(value);
576 }
577 
578 void
set_env(std::string const & key,std::string const & value)579 Arguments::set_env(std::string const &key, std::string const &value)
580 {
581   // perform overwrite for now
582   _data_map[key]._env_value = value;
583 }
584 
585 void
show_all_configuration() const586 Arguments::show_all_configuration() const
587 {
588   for (const auto &it : _data_map) {
589     std::cout << "name: " + it.first << std::endl;
590     std::string msg;
591     msg = "args value:";
592     for (const auto &it_data : it.second._values) {
593       msg += " " + it_data;
594     }
595     std::cout << msg << std::endl;
596     std::cout << "env value: " + it.second._env_value << std::endl << std::endl;
597   }
598 }
599 
600 // invoke the function with the args
601 void
invoke()602 Arguments::invoke()
603 {
604   if (_action) {
605     // call the std::function
606     _action();
607   } else {
608     throw std::runtime_error("no function to invoke");
609   }
610 }
611 
612 bool
has_action() const613 Arguments::has_action() const
614 {
615   return _action != nullptr;
616 }
617 
618 //=========================== ArgumentData class ================================
619 
620 std::string const &
env() const621 ArgumentData::env() const noexcept
622 {
623   return _env_value;
624 }
625 
626 std::string const &
at(unsigned index) const627 ArgumentData::at(unsigned index) const
628 {
629   if (index >= _values.size()) {
630     throw std::out_of_range("argument not found at index: " + std::to_string(index));
631   }
632   return _values.at(index);
633 }
634 
635 std::string const &
value() const636 ArgumentData::value() const noexcept
637 {
638   if (_values.empty()) {
639     // To prevent compiler warning
640     static const std::string empty = "";
641     return empty;
642   }
643   return _values.at(0);
644 }
645 
646 size_t
size() const647 ArgumentData::size() const noexcept
648 {
649   return _values.size();
650 }
651 
652 bool
empty() const653 ArgumentData::empty() const noexcept
654 {
655   return _values.empty() && _env_value.empty();
656 }
657 
658 AP_StrVec::const_iterator
begin() const659 ArgumentData::begin() const noexcept
660 {
661   return _values.begin();
662 }
663 
664 AP_StrVec::const_iterator
end() const665 ArgumentData::end() const noexcept
666 {
667   return _values.end();
668 }
669 
670 } // namespace ts
671