crosh/src/util.rs - third_party/platform2 - Git at Google

 // Copyright 2019 The ChromiumOS Authors
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 // Provides helper functions used by handler implementations of crosh commands.

 use rand::distributions::Alphanumeric;
 use rand::{thread_rng, Rng};
 use std::env;
 use std::error;
 use std::fmt::{self, Display};
 use std::fs::read_to_string;
 use std::io::stdin;
 use std::io::stdout;
 use std::io::Read;
 use std::io::Write;
 use std::process::{Command, Stdio};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::time::Duration;

 use chrono::Local;
 use dbus::blocking::Connection;
 use libc::{c_int, epoll_event};
 use libchromeos::chromeos;
 use libchromeos::signal::{clear_signal_handler, register_signal_handler};
 use log::error;
 use nix::sys::signal::Signal;
 use system_api::client::OrgChromiumSessionManagerInterface;

 // 25 seconds is the default timeout for dbus-send.
 pub const DEFAULT_DBUS_TIMEOUT: Duration = Duration::from_secs(25);
 // Path to update_engine_client.
 pub const UPDATE_ENGINE: &str = "/usr/bin/update_engine_client";

 const CROS_USER_ID_HASH: &str = "CROS_USER_ID_HASH";

 // 30 seconds is the default timeout for epoll.
 const DEFAULT_EPOLL_TIMEOUT: i32 = 30000;
 const MAX_EPOLL_EVENTS: i32 = 1;

 static INCLUDE_DEV: AtomicBool = AtomicBool::new(false);
 static INCLUDE_USB: AtomicBool = AtomicBool::new(false);

 #[derive(Debug)]
 pub enum Error {
     DbusChromeFeaturesService(dbus::Error, String),
     DbusConnection(dbus::Error),
     DbusGetUserIdHash(chromeos::Error),
     DbusGuestSessionActive(dbus::Error),
     EpollAdd(i32),
     EpollCreate(i32),
     NoMatchFound,
     WrappedError(String),
 }

 pub type Result<T> = std::result::Result<T, Error>;

 impl Display for Error {
     #[remain::check]
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         use self::Error::*;

         #[sorted]
         match self {
             DbusChromeFeaturesService(err, m) => write!(f, "failed to call '{}': {}", m, err),
             DbusConnection(err) => write!(f, "failed to connect to D-Bus: {}", err),
             DbusGetUserIdHash(err) => {
                 write!(f, "failed to get user-id hash over to D-Bus: {:?}", err)
             }
             DbusGuestSessionActive(err) => {
                 write!(
                     f,
                     "failed to check whether guest session is active over D-Bus: {:?}",
                     err
                 )
             }
             EpollAdd(err) => write!(f, "failed to add to epoll: {}", err),
             EpollCreate(err) => write!(f, "failed to create epoll: {}", err),
             NoMatchFound => write!(f, "No match found."),
             WrappedError(err) => write!(f, "{}", err),
         }
     }
 }

 impl<T: error::Error> From<T> for Error {
     fn from(err: T) -> Self {
         Error::WrappedError(format!("{:?}", err))
     }
 }

 // Return the user ID hash from the environment. If it is not available, fetch it from session
 // manager and set the environment variable.
 pub fn get_user_id_hash() -> Result<String> {
     if let Ok(lookup) = env::var(CROS_USER_ID_HASH) {
         return Ok(lookup);
     }

     let user_id_hash = chromeos::get_user_id_hash().map_err(|err| {
         error!("ERROR: D-Bus call failed: {}", err);
         Error::DbusGetUserIdHash(err)
     })?;

     env::set_var(CROS_USER_ID_HASH, &user_id_hash);
     Ok(user_id_hash)
 }

 // Return the output file path for the given output type in user's Downloads directory.
 pub fn generate_output_file_path(output_type: &str, file_extension: &str) -> Result<String> {
     let date = Local::now();
     let formatted_date = date.format("%Y-%m-%d_%H.%M.%S");
     let user_id_hash = get_user_id_hash()?;
     let random_string: String = thread_rng().sample_iter(&Alphanumeric).take(6).collect();
     Ok(format!(
         "/home/user/{}/MyFiles/Downloads/{}_{}_{}.{}",
         user_id_hash, output_type, formatted_date, random_string, file_extension
     ))
 }

 pub fn is_chrome_feature_enabled(method_name: &str) -> Result<bool> {
     let user_id_hash = get_user_id_hash()?;

     let connection = Connection::new_system().map_err(|err| {
         error!("ERROR: Failed to get D-Bus connection: {}", err);
         Error::DbusConnection(err)
     })?;

     let proxy = connection.with_proxy(
         "org.chromium.ChromeFeaturesService",
         "/org/chromium/ChromeFeaturesService",
         DEFAULT_DBUS_TIMEOUT,
     );

     let (reply,): (bool,) = proxy
         .method_call(
             "org.chromium.ChromeFeaturesServiceInterface",
             method_name,
             (user_id_hash,),
         )
         .map_err(|err| {
             error!("ERROR: D-Bus method call failed: {}", err);
             Error::DbusChromeFeaturesService(err, method_name.to_string())
         })?;

     Ok(reply)
 }

 pub fn is_no_new_privs_set() -> bool {
     // Safe because this retrieves a value without side effects.
     unsafe { libc::prctl(libc::PR_GET_NO_NEW_PRIVS, 0, 0, 0, 0) > 0 }
 }

 pub fn is_removable() -> Result<bool> {
     let dev = root_dev()?;
     let dev = blockdev_basename(&dev).ok_or(Error::NoMatchFound)?;
     match read_to_string(format!("/sys/block/{}/removable", dev)) {
         Ok(contents) => Ok(contents.trim() == "1"),
         Err(err) => Err(err.into()),
     }
 }

 pub fn is_consumer_device() -> Result<bool> {
     let output = Command::new("/usr/sbin/device_management_client")
         .arg("--action=enterprise_owned_get_status")
         .output()?;

     let stdout = std::str::from_utf8(&output.stdout)
         .unwrap()
         .trim()
         .to_string();
     match output.status.code() {
         Some(0) if stdout == "EnterpriseOwnedGetStatus(): 0" => Ok(true),
         Some(0) if stdout == "EnterpriseOwnedGetStatus(): 1" => Ok(false),
         Some(0) => Err(Error::WrappedError(stdout)),
         Some(v) => Err(Error::WrappedError(format!("unexpected exit code: {}", v))),
         None => Err(Error::WrappedError("failed to get exit code".to_string())),
     }
 }

 pub fn set_dev_commands_included(value: bool) {
     INCLUDE_DEV.store(value, Ordering::Release);
 }

 pub fn set_usb_commands_included(value: bool) {
     INCLUDE_USB.store(value, Ordering::Release);
 }

 pub fn dev_commands_included() -> bool {
     INCLUDE_DEV.load(Ordering::Acquire)
 }

 pub fn usb_commands_included() -> bool {
     INCLUDE_USB.load(Ordering::Acquire)
 }

 /// # Safety
 /// handler needs to be async safe.
 pub unsafe fn set_signal_handlers(signums: &[Signal], handler: extern "C" fn(c_int)) {
     for signum in signums {
         // Safe as long as handler is async safe.
         if unsafe { register_signal_handler(*signum, handler) }.is_err() {
             error!("sigaction failed for {}", signum);
         }
     }
 }

 pub fn clear_signal_handlers(signums: &[Signal]) {
     for signum in signums {
         if clear_signal_handler(*signum).is_err() {
             error!("sigaction failed for {}", signum);
         }
     }
 }

 pub fn is_guest_session_active() -> Result<bool> {
     let connection = Connection::new_system().map_err(|err| {
         error!("ERROR: Failed to get D-Bus connection: {}", err);
         Error::DbusConnection(err)
     })?;

     let conn_path = connection.with_proxy(
         "org.chromium.SessionManager",
         "/org/chromium/SessionManager",
         DEFAULT_DBUS_TIMEOUT,
     );

     let guest_session_active = conn_path.is_guest_session_active().map_err(|err| {
         println!("ERROR: Got unexpected result: {}", err);
         Error::DbusGuestSessionActive(err)
     })?;

     Ok(guest_session_active)
 }

 fn root_dev() -> Result<String> {
     let mut child = Command::new("rootdev")
         .arg("-s")
         .stdin(Stdio::null())
         .stdout(Stdio::piped())
         .spawn()?;

     let mut result = String::new();
     child.stdout.take().unwrap().read_to_string(&mut result)?;
     child.wait()?;

     Ok(result.trim().to_string())
 }

 /// Given a path to a block device with partition number, return its name.
 ///
 /// For example: "/dev/nvme0n1p1" -> "nvme0n1"
 fn blockdev_basename(mut dev: &str) -> Option<&str> {
     dev = dev.strip_prefix("/dev/")?;

     // The device name should be a single path component.
     if dev.contains('/') {
         return None;
     }

     // Remove partition number.
     if let Some((base, part)) = dev.rsplit_once('p') {
         // nvme0n1p3 -> nvme0n1
         if part.chars().all(|c: char| c.is_ascii_digit()) {
             return Some(base);
         }
     }

     // sda1 -> sda
     Some(dev.trim_end_matches(|c: char| c.is_ascii_digit()))
 }

 /// Print 'msg' followed by a [y/N] prompt and test the user input. Return true for 'y' or 'Y'.
 pub fn prompt_for_yes(msg: &str) -> bool {
     print!("{} [y/N] ", msg);
     stdout().flush().ok();

     let mut response = String::new();
     stdin().read_line(&mut response).ok();
     matches!(response.as_str(), "y\n" | "Y\n")
 }

 // Creates a new epoll instance watching the provided fd.
 pub fn add_epoll_for_fd(fd: i32) -> Result<i32> {
     // Safe because this will create a new epoll instance. Then we check the
     // return value.
     let epoll_fd = unsafe { libc::epoll_create1(0) };
     if epoll_fd < 0 {
         error!(
             "ERROR: Got unexpected result from epoll_create1(): {}",
             epoll_fd
         );
         return Err(Error::EpollCreate(epoll_fd));
     }

     let mut evt = epoll_event { events: 0, u64: 0 };
     // Safe because this only adds an entry to the epoll interest list. Then we
     // check the return value.
     let add_result = unsafe { libc::epoll_ctl(epoll_fd, libc::EPOLL_CTL_ADD, fd, &mut evt) };
     if add_result != 0 {
         error!("ERROR: Got unexpected epoll result: {}", add_result);
         return Err(Error::EpollAdd(add_result));
     }

     Ok(epoll_fd)
 }

 // Waits for an epoll event on the provided fd.
 //
 // The return value of the underlying syscall is returned.
 pub fn epoll_wait(epoll_fd: i32) -> i32 {
     let mut evt = epoll_event { events: 0, u64: 0 };
     // Safe because we give an epoll_event array that can hold MAX_EPOLL_EVENTS.
     unsafe { libc::epoll_wait(epoll_fd, &mut evt, MAX_EPOLL_EVENTS, DEFAULT_EPOLL_TIMEOUT) }
 }

 #[cfg(test)]
 mod tests {
     use regex::Regex;

     use super::*;

     #[test]
     fn test_generate_output_file_path() {
         // Set the user id hash env variable to a random value because it is necessary for output path generation.
         env::set_var(CROS_USER_ID_HASH, "useridhashfortesting");
         let expected_path_re = Regex::new(concat!(
             r"^/home/user/.+/MyFiles/Downloads/",
             r"packet_capture_\d{4}-\d{2}-\d{2}_\d{2}.\d{2}.\d{2}_.{6}\.pcap$"
         ))
         .unwrap();
         let result_output_path = generate_output_file_path("packet_capture", "pcap").unwrap();
         assert!(expected_path_re.is_match(&result_output_path));
     }

     #[test]
     fn blockdev_basename_valid() {
         assert_eq!(blockdev_basename("/dev/mmcblk0p3"), Some("mmcblk0"));
         assert_eq!(blockdev_basename("/dev/nvme0n1p1"), Some("nvme0n1"));
         assert_eq!(blockdev_basename("/dev/sda1"), Some("sda"));
     }

     #[test]
     fn blockdev_basename_invalid() {
         assert!(blockdev_basename("abcdef").is_none());
         assert!(blockdev_basename("/dev/abc/def").is_none());
     }
 }
	// Copyright 2019 The ChromiumOS Authors
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	// Provides helper functions used by handler implementations of crosh commands.

	use rand::distributions::Alphanumeric;
	use rand::{thread_rng, Rng};
	use std::env;
	use std::error;
	use std::fmt::{self, Display};
	use std::fs::read_to_string;
	use std::io::stdin;
	use std::io::stdout;
	use std::io::Read;
	use std::io::Write;
	use std::process::{Command, Stdio};
	use std::sync::atomic::{AtomicBool, Ordering};
	use std::time::Duration;

	use chrono::Local;
	use dbus::blocking::Connection;
	use libc::{c_int, epoll_event};
	use libchromeos::chromeos;
	use libchromeos::signal::{clear_signal_handler, register_signal_handler};
	use log::error;
	use nix::sys::signal::Signal;
	use system_api::client::OrgChromiumSessionManagerInterface;

	// 25 seconds is the default timeout for dbus-send.
	pub const DEFAULT_DBUS_TIMEOUT: Duration = Duration::from_secs(25);
	// Path to update_engine_client.
	pub const UPDATE_ENGINE: &str = "/usr/bin/update_engine_client";

	const CROS_USER_ID_HASH: &str = "CROS_USER_ID_HASH";

	// 30 seconds is the default timeout for epoll.
	const DEFAULT_EPOLL_TIMEOUT: i32 = 30000;
	const MAX_EPOLL_EVENTS: i32 = 1;

	static INCLUDE_DEV: AtomicBool = AtomicBool::new(false);
	static INCLUDE_USB: AtomicBool = AtomicBool::new(false);

	#[derive(Debug)]
	pub enum Error {
	DbusChromeFeaturesService(dbus::Error, String),
	DbusConnection(dbus::Error),
	DbusGetUserIdHash(chromeos::Error),
	DbusGuestSessionActive(dbus::Error),
	EpollAdd(i32),
	EpollCreate(i32),
	NoMatchFound,
	WrappedError(String),
	}

	pub type Result<T> = std::result::Result<T, Error>;

	impl Display for Error {
	#[remain::check]
	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
	use self::Error::*;

	#[sorted]
	match self {
	DbusChromeFeaturesService(err, m) => write!(f, "failed to call '{}': {}", m, err),
	DbusConnection(err) => write!(f, "failed to connect to D-Bus: {}", err),
	DbusGetUserIdHash(err) => {
	write!(f, "failed to get user-id hash over to D-Bus: {:?}", err)
	}
	DbusGuestSessionActive(err) => {
	write!(
	f,
	"failed to check whether guest session is active over D-Bus: {:?}",
	err
	)
	}
	EpollAdd(err) => write!(f, "failed to add to epoll: {}", err),
	EpollCreate(err) => write!(f, "failed to create epoll: {}", err),
	NoMatchFound => write!(f, "No match found."),
	WrappedError(err) => write!(f, "{}", err),
	}
	}
	}

	impl<T: error::Error> From<T> for Error {
	fn from(err: T) -> Self {
	Error::WrappedError(format!("{:?}", err))
	}
	}

	// Return the user ID hash from the environment. If it is not available, fetch it from session
	// manager and set the environment variable.
	pub fn get_user_id_hash() -> Result<String> {
	if let Ok(lookup) = env::var(CROS_USER_ID_HASH) {
	return Ok(lookup);
	}

	let user_id_hash = chromeos::get_user_id_hash().map_err(\|err\| {
	error!("ERROR: D-Bus call failed: {}", err);
	Error::DbusGetUserIdHash(err)
	})?;

	env::set_var(CROS_USER_ID_HASH, &user_id_hash);
	Ok(user_id_hash)
	}

	// Return the output file path for the given output type in user's Downloads directory.
	pub fn generate_output_file_path(output_type: &str, file_extension: &str) -> Result<String> {
	let date = Local::now();
	let formatted_date = date.format("%Y-%m-%d_%H.%M.%S");
	let user_id_hash = get_user_id_hash()?;
	let random_string: String = thread_rng().sample_iter(&Alphanumeric).take(6).collect();
	Ok(format!(
	"/home/user/{}/MyFiles/Downloads/{}_{}_{}.{}",
	user_id_hash, output_type, formatted_date, random_string, file_extension
	))
	}

	pub fn is_chrome_feature_enabled(method_name: &str) -> Result<bool> {
	let user_id_hash = get_user_id_hash()?;

	let connection = Connection::new_system().map_err(\|err\| {
	error!("ERROR: Failed to get D-Bus connection: {}", err);
	Error::DbusConnection(err)
	})?;

	let proxy = connection.with_proxy(
	"org.chromium.ChromeFeaturesService",
	"/org/chromium/ChromeFeaturesService",
	DEFAULT_DBUS_TIMEOUT,
	);

	let (reply,): (bool,) = proxy
	.method_call(
	"org.chromium.ChromeFeaturesServiceInterface",
	method_name,
	(user_id_hash,),
	)
	.map_err(\|err\| {
	error!("ERROR: D-Bus method call failed: {}", err);
	Error::DbusChromeFeaturesService(err, method_name.to_string())
	})?;

	Ok(reply)
	}

	pub fn is_no_new_privs_set() -> bool {
	// Safe because this retrieves a value without side effects.
	unsafe { libc::prctl(libc::PR_GET_NO_NEW_PRIVS, 0, 0, 0, 0) > 0 }
	}

	pub fn is_removable() -> Result<bool> {
	let dev = root_dev()?;
	let dev = blockdev_basename(&dev).ok_or(Error::NoMatchFound)?;
	match read_to_string(format!("/sys/block/{}/removable", dev)) {
	Ok(contents) => Ok(contents.trim() == "1"),
	Err(err) => Err(err.into()),
	}
	}

	pub fn is_consumer_device() -> Result<bool> {
	let output = Command::new("/usr/sbin/device_management_client")
	.arg("--action=enterprise_owned_get_status")
	.output()?;

	let stdout = std::str::from_utf8(&output.stdout)
	.unwrap()
	.trim()
	.to_string();
	match output.status.code() {
	Some(0) if stdout == "EnterpriseOwnedGetStatus(): 0" => Ok(true),
	Some(0) if stdout == "EnterpriseOwnedGetStatus(): 1" => Ok(false),
	Some(0) => Err(Error::WrappedError(stdout)),
	Some(v) => Err(Error::WrappedError(format!("unexpected exit code: {}", v))),
	None => Err(Error::WrappedError("failed to get exit code".to_string())),
	}
	}

	pub fn set_dev_commands_included(value: bool) {
	INCLUDE_DEV.store(value, Ordering::Release);
	}

	pub fn set_usb_commands_included(value: bool) {
	INCLUDE_USB.store(value, Ordering::Release);
	}

	pub fn dev_commands_included() -> bool {
	INCLUDE_DEV.load(Ordering::Acquire)
	}

	pub fn usb_commands_included() -> bool {
	INCLUDE_USB.load(Ordering::Acquire)
	}

	/// # Safety
	/// handler needs to be async safe.
	pub unsafe fn set_signal_handlers(signums: &[Signal], handler: extern "C" fn(c_int)) {
	for signum in signums {
	// Safe as long as handler is async safe.
	if unsafe { register_signal_handler(*signum, handler) }.is_err() {
	error!("sigaction failed for {}", signum);
	}
	}
	}

	pub fn clear_signal_handlers(signums: &[Signal]) {
	for signum in signums {
	if clear_signal_handler(*signum).is_err() {
	error!("sigaction failed for {}", signum);
	}
	}
	}

	pub fn is_guest_session_active() -> Result<bool> {
	let connection = Connection::new_system().map_err(\|err\| {
	error!("ERROR: Failed to get D-Bus connection: {}", err);
	Error::DbusConnection(err)
	})?;

	let conn_path = connection.with_proxy(
	"org.chromium.SessionManager",
	"/org/chromium/SessionManager",
	DEFAULT_DBUS_TIMEOUT,
	);

	let guest_session_active = conn_path.is_guest_session_active().map_err(\|err\| {
	println!("ERROR: Got unexpected result: {}", err);
	Error::DbusGuestSessionActive(err)
	})?;

	Ok(guest_session_active)
	}

	fn root_dev() -> Result<String> {
	let mut child = Command::new("rootdev")
	.arg("-s")
	.stdin(Stdio::null())
	.stdout(Stdio::piped())
	.spawn()?;

	let mut result = String::new();
	child.stdout.take().unwrap().read_to_string(&mut result)?;
	child.wait()?;

	Ok(result.trim().to_string())
	}

	/// Given a path to a block device with partition number, return its name.
	///
	/// For example: "/dev/nvme0n1p1" -> "nvme0n1"
	fn blockdev_basename(mut dev: &str) -> Option<&str> {
	dev = dev.strip_prefix("/dev/")?;

	// The device name should be a single path component.
	if dev.contains('/') {
	return None;
	}

	// Remove partition number.
	if let Some((base, part)) = dev.rsplit_once('p') {
	// nvme0n1p3 -> nvme0n1
	if part.chars().all(\|c: char\| c.is_ascii_digit()) {
	return Some(base);
	}
	}

	// sda1 -> sda
	Some(dev.trim_end_matches(\|c: char\| c.is_ascii_digit()))
	}

	/// Print 'msg' followed by a [y/N] prompt and test the user input. Return true for 'y' or 'Y'.
	pub fn prompt_for_yes(msg: &str) -> bool {
	print!("{} [y/N] ", msg);
	stdout().flush().ok();

	let mut response = String::new();
	stdin().read_line(&mut response).ok();
	matches!(response.as_str(), "y\n" \| "Y\n")
	}

	// Creates a new epoll instance watching the provided fd.
	pub fn add_epoll_for_fd(fd: i32) -> Result<i32> {
	// Safe because this will create a new epoll instance. Then we check the
	// return value.
	let epoll_fd = unsafe { libc::epoll_create1(0) };
	if epoll_fd < 0 {
	error!(
	"ERROR: Got unexpected result from epoll_create1(): {}",
	epoll_fd
	);
	return Err(Error::EpollCreate(epoll_fd));
	}

	let mut evt = epoll_event { events: 0, u64: 0 };
	// Safe because this only adds an entry to the epoll interest list. Then we
	// check the return value.
	let add_result = unsafe { libc::epoll_ctl(epoll_fd, libc::EPOLL_CTL_ADD, fd, &mut evt) };
	if add_result != 0 {
	error!("ERROR: Got unexpected epoll result: {}", add_result);
	return Err(Error::EpollAdd(add_result));
	}

	Ok(epoll_fd)
	}

	// Waits for an epoll event on the provided fd.
	//
	// The return value of the underlying syscall is returned.
	pub fn epoll_wait(epoll_fd: i32) -> i32 {
	let mut evt = epoll_event { events: 0, u64: 0 };
	// Safe because we give an epoll_event array that can hold MAX_EPOLL_EVENTS.
	unsafe { libc::epoll_wait(epoll_fd, &mut evt, MAX_EPOLL_EVENTS, DEFAULT_EPOLL_TIMEOUT) }
	}

	#[cfg(test)]
	mod tests {
	use regex::Regex;

	use super::*;

	#[test]
	fn test_generate_output_file_path() {
	// Set the user id hash env variable to a random value because it is necessary for output path generation.
	env::set_var(CROS_USER_ID_HASH, "useridhashfortesting");
	let expected_path_re = Regex::new(concat!(
	r"^/home/user/.+/MyFiles/Downloads/",
	r"packet_capture_\d{4}-\d{2}-\d{2}_\d{2}.\d{2}.\d{2}_.{6}\.pcap$"
	))
	.unwrap();
	let result_output_path = generate_output_file_path("packet_capture", "pcap").unwrap();
	assert!(expected_path_re.is_match(&result_output_path));
	}

	#[test]
	fn blockdev_basename_valid() {
	assert_eq!(blockdev_basename("/dev/mmcblk0p3"), Some("mmcblk0"));
	assert_eq!(blockdev_basename("/dev/nvme0n1p1"), Some("nvme0n1"));
	assert_eq!(blockdev_basename("/dev/sda1"), Some("sda"));
	}

	#[test]
	fn blockdev_basename_invalid() {
	assert!(blockdev_basename("abcdef").is_none());
	assert!(blockdev_basename("/dev/abc/def").is_none());
	}
	}