client.cc (14286B)
1 // Copyright 2022 The Chromium Authors. 2 // Use of this source code is governed by a BSD-style license that can be 3 // found in the LICENSE file. 4 5 #include <time.h> 6 7 #include <iostream> 8 #include <memory> 9 #include <mutex> 10 #include <sstream> 11 #include <string> 12 #include <thread> 13 #include <vector> 14 15 #include "content_analysis/sdk/analysis_client.h" 16 #include "demo/atomic_output.h" 17 18 using content_analysis::sdk::Client; 19 using content_analysis::sdk::ContentAnalysisRequest; 20 using content_analysis::sdk::ContentAnalysisResponse; 21 using content_analysis::sdk::ContentAnalysisAcknowledgement; 22 23 // Different paths are used depending on whether this agent should run as a 24 // use specific agent or not. These values are chosen to match the test 25 // values in chrome browser. 26 constexpr char kPathUser[] = "path_user"; 27 constexpr char kPathSystem[] = "brcm_chrm_cas"; 28 29 // Global app config. 30 std::string path = kPathSystem; 31 bool user_specific = false; 32 bool group = false; 33 std::unique_ptr<Client> client; 34 35 // Paramters used to build the request. 36 content_analysis::sdk::AnalysisConnector connector = 37 content_analysis::sdk::FILE_ATTACHED; 38 time_t request_token_number = time(nullptr); 39 std::string request_token; 40 std::string tag = "dlp"; 41 bool threaded = false; 42 std::string digest = "sha256-123456"; 43 std::string url = "https://upload.example.com"; 44 std::string email = "me@example.com"; 45 std::string machine_user = "DOMAIN\\me"; 46 std::vector<std::string> datas; 47 48 // When grouping, remember the tokens of all requests/responses in order to 49 // acknowledge them all with the same final action. 50 // 51 // This global state. It may be access from multiple thread so must be 52 // accessed from a critical section. 53 std::mutex global_mutex; 54 ContentAnalysisAcknowledgement::FinalAction global_final_action = 55 ContentAnalysisAcknowledgement::ALLOW; 56 std::vector<std::string> request_tokens; 57 58 // Command line parameters. 59 constexpr const char* kArgConnector = "--connector="; 60 constexpr const char* kArgDigest = "--digest="; 61 constexpr const char* kArgEmail = "--email="; 62 constexpr const char* kArgGroup = "--group"; 63 constexpr const char* kArgMachineUser = "--machine-user="; 64 constexpr const char* kArgPath = "--path="; 65 constexpr const char* kArgRequestToken = "--request-token="; 66 constexpr const char* kArgTag = "--tag="; 67 constexpr const char* kArgThreaded = "--threaded"; 68 constexpr const char* kArgUrl = "--url="; 69 constexpr const char* kArgUserSpecific = "--user"; 70 constexpr const char* kArgHelp = "--help"; 71 72 bool ParseCommandLine(int argc, char* argv[]) { 73 for (int i = 1; i < argc; ++i) { 74 const std::string arg = argv[i]; 75 if (arg.find(kArgConnector) == 0) { 76 std::string connector_str = arg.substr(strlen(kArgConnector)); 77 if (connector_str == "download") { 78 connector = content_analysis::sdk::FILE_DOWNLOADED; 79 } else if (connector_str == "attach") { 80 connector = content_analysis::sdk::FILE_ATTACHED; 81 } else if (connector_str == "bulk-data-entry") { 82 connector = content_analysis::sdk::BULK_DATA_ENTRY; 83 } else if (connector_str == "print") { 84 connector = content_analysis::sdk::PRINT; 85 } else if (connector_str == "file-transfer") { 86 connector = content_analysis::sdk::FILE_TRANSFER; 87 } else { 88 std::cout << "[Demo] Incorrect command line arg: " << arg << std::endl; 89 return false; 90 } 91 } else if (arg.find(kArgRequestToken) == 0) { 92 request_token = arg.substr(strlen(kArgRequestToken)); 93 } else if (arg.find(kArgTag) == 0) { 94 tag = arg.substr(strlen(kArgTag)); 95 } else if (arg.find(kArgThreaded) == 0) { 96 threaded = true; 97 } else if (arg.find(kArgDigest) == 0) { 98 digest = arg.substr(strlen(kArgDigest)); 99 } else if (arg.find(kArgUrl) == 0) { 100 url = arg.substr(strlen(kArgUrl)); 101 } else if (arg.find(kArgMachineUser) == 0) { 102 machine_user = arg.substr(strlen(kArgMachineUser)); 103 } else if (arg.find(kArgEmail) == 0) { 104 email = arg.substr(strlen(kArgEmail)); 105 } else if (arg.find(kArgPath) == 0) { 106 path = arg.substr(strlen(kArgPath)); 107 } else if (arg.find(kArgUserSpecific) == 0) { 108 // If kArgPath was already used, abort. 109 if (path != kPathSystem) { 110 std::cout << std::endl << "ERROR: use --path=<path> after --user"; 111 return false; 112 } 113 path = kPathUser; 114 user_specific = true; 115 } else if (arg.find(kArgGroup) == 0) { 116 group = true; 117 } else if (arg.find(kArgHelp) == 0) { 118 return false; 119 } else { 120 datas.push_back(arg); 121 } 122 } 123 124 return true; 125 } 126 127 void PrintHelp() { 128 std::cout 129 << std::endl << std::endl 130 << "Usage: client [OPTIONS] [@]content_or_file ..." << std::endl 131 << "A simple client to send content analysis requests to a running agent." << std::endl 132 << "Without @ the content to analyze is the argument itself." << std::endl 133 << "Otherwise the content is read from a file called 'content_or_file'." << std::endl 134 << "Multiple [@]content_or_file arguments may be specified, each generates one request." << std::endl 135 << std::endl << "Options:" << std::endl 136 << kArgConnector << "<connector> : one of 'download', 'attach' (default), 'bulk-data-entry', 'print', or 'file-transfer'" << std::endl 137 << kArgRequestToken << "<unique-token> : defaults to 'req-<number>' which auto increments" << std::endl 138 << kArgTag << "<tag> : defaults to 'dlp'" << std::endl 139 << kArgThreaded << " : handled multiple requests using threads" << std::endl 140 << kArgUrl << "<url> : defaults to 'https://upload.example.com'" << std::endl 141 << kArgMachineUser << "<machine-user> : defaults to 'DOMAIN\\me'" << std::endl 142 << kArgEmail << "<email> : defaults to 'me@example.com'" << std::endl 143 << kArgPath << " <path> : Used the specified path instead of default. Must come after --user." << std::endl 144 << kArgUserSpecific << " : Connects to an OS user specific agent" << std::endl 145 << kArgDigest << "<digest> : defaults to 'sha256-123456'" << std::endl 146 << kArgGroup << " : Generate the same final action for all requests" << std::endl 147 << kArgHelp << " : prints this help message" << std::endl; 148 } 149 150 std::string GenerateRequestToken() { 151 std::stringstream stm; 152 stm << "req-" << request_token_number++; 153 return stm.str(); 154 } 155 156 ContentAnalysisRequest BuildRequest(const std::string& data) { 157 std::string filepath; 158 std::string filename; 159 if (data[0] == '@') { 160 filepath = data.substr(1); 161 filename = filepath.substr(filepath.find_last_of("/\\") + 1); 162 } 163 164 ContentAnalysisRequest request; 165 166 // Set request to expire 5 minutes into the future. 167 request.set_expires_at(time(nullptr) + 5 * 60); 168 request.set_analysis_connector(connector); 169 request.set_request_token(!request_token.empty() 170 ? request_token : GenerateRequestToken()); 171 *request.add_tags() = tag; 172 173 auto request_data = request.mutable_request_data(); 174 request_data->set_url(url); 175 request_data->set_email(email); 176 request_data->set_digest(digest); 177 if (!filename.empty()) { 178 request_data->set_filename(filename); 179 } 180 181 auto client_metadata = request.mutable_client_metadata(); 182 auto browser = client_metadata->mutable_browser(); 183 browser->set_machine_user(machine_user); 184 185 if (!filepath.empty()) { 186 request.set_file_path(filepath); 187 } else if (!data.empty()) { 188 request.set_text_content(data); 189 } else { 190 std::cout << "[Demo] Specify text content or a file path." << std::endl; 191 PrintHelp(); 192 exit(1); 193 } 194 195 return request; 196 } 197 198 // Gets the most severe action within the result. 199 ContentAnalysisResponse::Result::TriggeredRule::Action 200 GetActionFromResult(const ContentAnalysisResponse::Result& result) { 201 auto action = 202 ContentAnalysisResponse::Result::TriggeredRule::ACTION_UNSPECIFIED; 203 for (auto rule : result.triggered_rules()) { 204 if (rule.has_action() && rule.action() > action) 205 action = rule.action(); 206 } 207 return action; 208 } 209 210 // Gets the most severe action within all the the results of a response. 211 ContentAnalysisResponse::Result::TriggeredRule::Action 212 GetActionFromResponse(const ContentAnalysisResponse& response) { 213 auto action = 214 ContentAnalysisResponse::Result::TriggeredRule::ACTION_UNSPECIFIED; 215 for (auto result : response.results()) { 216 auto action2 = GetActionFromResult(result); 217 if (action2 > action) 218 action = action2; 219 } 220 return action; 221 } 222 223 void DumpResponse( 224 std::stringstream& stream, 225 const ContentAnalysisResponse& response) { 226 for (auto result : response.results()) { 227 auto tag = result.has_tag() ? result.tag() : "<no-tag>"; 228 229 auto status = result.has_status() 230 ? result.status() 231 : ContentAnalysisResponse::Result::STATUS_UNKNOWN; 232 std::string status_str; 233 switch (status) { 234 case ContentAnalysisResponse::Result::STATUS_UNKNOWN: 235 status_str = "Unknown"; 236 break; 237 case ContentAnalysisResponse::Result::SUCCESS: 238 status_str = "Success"; 239 break; 240 case ContentAnalysisResponse::Result::FAILURE: 241 status_str = "Failure"; 242 break; 243 default: 244 status_str = "<Uknown>"; 245 break; 246 } 247 248 auto action = GetActionFromResult(result); 249 std::string action_str; 250 switch (action) { 251 case ContentAnalysisResponse::Result::TriggeredRule::ACTION_UNSPECIFIED: 252 action_str = "allowed"; 253 break; 254 case ContentAnalysisResponse::Result::TriggeredRule::REPORT_ONLY: 255 action_str = "reported only"; 256 break; 257 case ContentAnalysisResponse::Result::TriggeredRule::WARN: 258 action_str = "warned"; 259 break; 260 case ContentAnalysisResponse::Result::TriggeredRule::BLOCK: 261 action_str = "blocked"; 262 break; 263 } 264 265 time_t now = time(nullptr); 266 stream << "[Demo] Request " << response.request_token() << " is " << action_str 267 << " after " << tag 268 << " analysis, status=" << status_str 269 << " at " << ctime(&now); 270 } 271 } 272 273 ContentAnalysisAcknowledgement BuildAcknowledgement( 274 const std::string& request_token, 275 ContentAnalysisAcknowledgement::FinalAction final_action) { 276 ContentAnalysisAcknowledgement ack; 277 ack.set_request_token(request_token); 278 ack.set_status(ContentAnalysisAcknowledgement::SUCCESS); 279 ack.set_final_action(final_action); 280 return ack; 281 } 282 283 void HandleRequest(const ContentAnalysisRequest& request) { 284 AtomicCout aout; 285 ContentAnalysisResponse response; 286 int err = client->Send(request, &response); 287 if (err != 0) { 288 aout.stream() << "[Demo] Error sending request " << request.request_token() 289 << std::endl; 290 } else if (response.results_size() == 0) { 291 aout.stream() << "[Demo] Response " << request.request_token() << " is missing a result" 292 << std::endl; 293 } else { 294 DumpResponse(aout.stream(), response); 295 296 auto final_action = ContentAnalysisAcknowledgement::ALLOW; 297 switch (GetActionFromResponse(response)) { 298 case ContentAnalysisResponse::Result::TriggeredRule::ACTION_UNSPECIFIED: 299 break; 300 case ContentAnalysisResponse::Result::TriggeredRule::REPORT_ONLY: 301 final_action = ContentAnalysisAcknowledgement::REPORT_ONLY; 302 break; 303 case ContentAnalysisResponse::Result::TriggeredRule::WARN: 304 final_action = ContentAnalysisAcknowledgement::WARN; 305 break; 306 case ContentAnalysisResponse::Result::TriggeredRule::BLOCK: 307 final_action = ContentAnalysisAcknowledgement::BLOCK; 308 break; 309 } 310 311 // If grouping, remember the request's token in order to ack the response 312 // later. 313 if (group) { 314 std::unique_lock<std::mutex> lock(global_mutex); 315 request_tokens.push_back(request.request_token()); 316 if (final_action > global_final_action) 317 global_final_action = final_action; 318 } else { 319 int err = client->Acknowledge( 320 BuildAcknowledgement(response.request_token(), final_action)); 321 if (err != 0) { 322 aout.stream() << "[Demo] Error sending ack " << request.request_token() 323 << std::endl; 324 } 325 } 326 } 327 } 328 329 void ProcessRequest(size_t i) { 330 auto request = BuildRequest(datas[i]); 331 332 { 333 AtomicCout aout; 334 aout.stream() << "[Demo] Sending request " << request.request_token() << std::endl; 335 } 336 337 HandleRequest(request); 338 } 339 340 int main(int argc, char* argv[]) { 341 if (!ParseCommandLine(argc, argv)) { 342 PrintHelp(); 343 return 1; 344 } 345 346 // Each client uses a unique name to identify itself with Google Chrome. 347 client = Client::Create({path, user_specific}); 348 if (!client) { 349 std::cout << "[Demo] Error starting client" << std::endl; 350 return 1; 351 }; 352 353 auto info = client->GetAgentInfo(); 354 std::cout << "Agent pid=" << info.pid 355 << " path=" << info.binary_path << std::endl; 356 357 if (threaded) { 358 std::vector<std::unique_ptr<std::thread>> threads; 359 for (int i = 0; i < datas.size(); ++i) { 360 AtomicCout aout; 361 aout.stream() << "Start thread " << i << std::endl; 362 threads.emplace_back(std::make_unique<std::thread>(ProcessRequest, i)); 363 } 364 365 // Make sure all threads have terminated. 366 for (auto& thread : threads) { 367 thread->join(); 368 } 369 } 370 else { 371 for (size_t i = 0; i < datas.size(); ++i) { 372 ProcessRequest(i); 373 } 374 } 375 // It's safe to access global state beyond this point without locking since 376 // all no more responses will be touching them. 377 378 if (group) { 379 std::cout << std::endl; 380 std::cout << "[Demo] Final action for all requests is "; 381 switch (global_final_action) { 382 // Google Chrome fails open, so if no action is specified that is the same 383 // as ALLOW. 384 case ContentAnalysisAcknowledgement::ACTION_UNSPECIFIED: 385 case ContentAnalysisAcknowledgement::ALLOW: 386 std::cout << "allowed"; 387 break; 388 case ContentAnalysisAcknowledgement::REPORT_ONLY: 389 std::cout << "reported only"; 390 break; 391 case ContentAnalysisAcknowledgement::WARN: 392 std::cout << "warned"; 393 break; 394 case ContentAnalysisAcknowledgement::BLOCK: 395 std::cout << "blocked"; 396 break; 397 } 398 std::cout << std::endl << std::endl; 399 400 for (auto token : request_tokens) { 401 std::cout << "[Demo] Sending group Ack" << std::endl; 402 int err = client->Acknowledge( 403 BuildAcknowledgement(token, global_final_action)); 404 if (err != 0) { 405 std::cout << "[Demo] Error sending ack for " << token << std::endl; 406 } 407 } 408 } 409 410 return 0; 411 };