1 /*
2  * Copyright (C) 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include "host/commands/snapshot_util_cvd/parse.h"
18 
19 #include <cstdlib>
20 #include <iostream>
21 #include <unordered_map>
22 
23 #include <android-base/parseint.h>
24 #include <android-base/strings.h>
25 
26 #include "common/libs/utils/contains.h"
27 #include "common/libs/utils/flag_parser.h"
28 #include "common/libs/utils/result.h"
29 #include "host/libs/config/cuttlefish_config.h"
30 
31 namespace cuttlefish {
32 namespace {
33 
34 constexpr char snapshot_cmd_help[] =
35     "Command to control regarding the snapshot operations: "
36     "suspend/resume/snapshot_take";
37 
38 constexpr char cleanup_snapshot_path_help[] =
39     "If true, snapshot_util_cvd will clean up the snapshot path on "
40     "failure of snapshot-taking";
41 
42 constexpr char wait_for_launcher_help[] =
43     "How many seconds to wait for the launcher to respond to the status "
44     "command. A value of zero means wait indefinitely.";
45 
46 constexpr char snapshot_path_help[] =
47     "Path to the directory the taken snapshot files are saved";
48 
SnapshotCmdFlag(std::string & value_buf)49 Flag SnapshotCmdFlag(std::string& value_buf) {
50   return GflagsCompatFlag("subcmd", value_buf).Help(snapshot_cmd_help);
51 }
52 
GetInt32Flag(const std::string & name,int & value_buf,const std::string & help_msg)53 Flag GetInt32Flag(const std::string& name, int& value_buf,
54                   const std::string& help_msg) {
55   return GflagsCompatFlag(name, value_buf).Help(help_msg);
56 }
57 
WaitForLauncherFlag(int & wait_for_launcher)58 Flag WaitForLauncherFlag(int& wait_for_launcher) {
59   return GetInt32Flag("wait_for_launcher", wait_for_launcher,
60                       wait_for_launcher_help);
61 }
62 
SnapshotPathFlag(std::string & path_buf)63 Flag SnapshotPathFlag(std::string& path_buf) {
64   return GflagsCompatFlag("snapshot_path", path_buf).Help(snapshot_path_help);
65 }
66 
CleanupSnapshotPathFlag(bool & cleanup)67 Flag CleanupSnapshotPathFlag(bool& cleanup) {
68   return GflagsCompatFlag("cleanup_snapshot_path", cleanup)
69       .Help(cleanup_snapshot_path_help);
70 }
71 
72 }  // namespace
73 
Parse(int argc,char ** argv)74 Result<Parsed> Parse(int argc, char** argv) {
75   auto args = ArgsToVec(argc, argv);
76   auto parsed = CF_EXPECT(Parse(args));
77   return parsed;
78 }
79 
ConvertToSnapshotCmd(const std::string & input)80 Result<SnapshotCmd> ConvertToSnapshotCmd(const std::string& input) {
81   std::unordered_map<std::string, SnapshotCmd> mapping{
82       {"suspend", SnapshotCmd::kSuspend},
83       {"resume", SnapshotCmd::kResume},
84       {"snapshot_take", SnapshotCmd::kSnapshotTake},
85       {"unknown", SnapshotCmd::kUnknown},
86   };
87   CF_EXPECT(Contains(mapping, input));
88   return mapping.at(input);
89 }
90 
InstanceNums()91 static Result<std::vector<int>> InstanceNums() {
92   CF_EXPECT(getenv("HOME") != nullptr, "\"HOME\" must be set properly.");
93   const auto* config = CuttlefishConfig::Get();
94   CF_EXPECT(config != nullptr, "CuttlefishConfig::Get() returned nullptr");
95 
96   const auto instances = config->Instances();
97   std::vector<int> instance_nums;
98   CF_EXPECT(!instances.empty(), "CuttlefishConfig has no instance in it.");
99   instance_nums.reserve(instances.size());
100   for (const auto& instance : instances) {
101     int id;
102     CF_EXPECTF(android::base::ParseInt(instance.id(), &id),
103                "Parsing filed for {}", id);
104     instance_nums.push_back(id);
105   }
106   return instance_nums;
107 }
108 
Parse(std::vector<std::string> & args)109 Result<Parsed> Parse(std::vector<std::string>& args) {
110   Parsed parsed{
111       .wait_for_launcher = 30,
112       .cleanup_snapshot_path = true,
113   };
114   std::vector<Flag> flags;
115   bool help_xml = false;
116   std::string snapshot_op("unknown");
117   std::string snapshot_path;
118   flags.push_back(SnapshotCmdFlag(snapshot_op));
119   flags.push_back(WaitForLauncherFlag(parsed.wait_for_launcher));
120   flags.push_back(SnapshotPathFlag(snapshot_path));
121   flags.push_back(CleanupSnapshotPathFlag(parsed.cleanup_snapshot_path));
122   flags.push_back(
123       GflagsCompatFlag("force", parsed.force)
124           .Help("If the snapshot path already exists, delete it first"));
125   flags.push_back(GflagsCompatFlag("auto_suspend", parsed.auto_suspend)
126                       .Help("Suspend/resume before/after taking the snapshot"));
127   flags.push_back(HelpFlag(flags));
128   flags.push_back(HelpXmlFlag(flags, std::cout, help_xml));
129   flags.push_back(UnexpectedArgumentGuard());
130   auto parse_res = ConsumeFlags(flags, args);
131   if (!help_xml && !parse_res.ok()) {
132     // Parse fails if helpxml is passed
133     CF_EXPECT(std::move(parse_res), "Flag parsing failed");
134   }
135   if (help_xml) {
136     std::exit(0);
137   }
138   parsed.cmd = CF_EXPECT(ConvertToSnapshotCmd(snapshot_op));
139   parsed.snapshot_path = snapshot_path;
140   parsed.instance_nums = CF_EXPECT(InstanceNums());
141   return parsed;
142 }
143 
operator <<(std::ostream & out,const SnapshotCmd & cmd)144 std::ostream& operator<<(std::ostream& out, const SnapshotCmd& cmd) {
145   switch (cmd) {
146     case SnapshotCmd::kUnknown:
147       out << "unknown";
148       break;
149     case SnapshotCmd::kSuspend:
150       out << "suspend";
151       break;
152     case SnapshotCmd::kResume:
153       out << "resume";
154       break;
155     case SnapshotCmd::kSnapshotTake:
156       out << "snapshot_take";
157       break;
158     default:
159       out << "unknown";
160       break;
161   }
162   return out;
163 }
164 
165 }  // namespace cuttlefish
166