Merge pull request #216 from rustcoreutils/ps

jgarzik · web-flow · commit 7c928a3f0448 · 2024-09-18T15:48:08.000+08:00
ps: support -a, -e, -d options
diff --git a/sys/ps.rs b/sys/ps.rs
@@ -1,3 +1,12 @@
+//
+// Copyright (c) 2024 Hemi Labs, Inc.
+//
+// This file is part of the posixutils-rs project covered under
+// the MIT License.  For the full license text, please see the LICENSE
+// file in the root directory of this project.
+// SPDX-License-Identifier: MIT
+//
+
 #[cfg(target_os = "macos")]
 mod psmacos;
 
@@ -23,26 +32,62 @@ struct Args {
     /// List all processes
     #[arg(short = 'A', long)]
     all: bool,
+
+    /// List all processes (alias for -A)
+    #[arg(short = 'e')]
+    all2: bool,
+
+    /// List all processes associated with terminals
+    #[arg(short = 'a', long)]
+    terminal_processes: bool,
+
+    /// Exclude session leaders
+    #[arg(short = 'd', long)]
+    exclude_session_leaders: bool,
 }
 
 fn main() {
-    let args = Args::parse();
-
-    if args.all {
-        match platform::list_processes() {
-            Ok(processes) => {
-                println!(
-                    "{:<5} {:<5} {:<5} {:<5} {}",
-                    "PID", "PPID", "UID", "GID", "COMMAND"
-                );
-                for proc in processes {
-                    println!(
-                        "{:<5} {:<5} {:<5} {:<5} {}",
-                        proc.pid, proc.ppid, proc.uid, proc.gid, proc.path
-                    );
-                }
-            }
-            Err(e) => eprintln!("Error: {}", e),
+    let mut args = Args::parse();
+    if args.all2 {
+        args.all = true;
+    }
+
+    let processes = match platform::list_processes() {
+        Ok(processes) => processes,
+        Err(e) => {
+            eprintln!("Error: {}", e);
+            return;
         }
+    };
+
+    // Filter processes based on arguments
+    let filtered_processes = if args.terminal_processes && args.exclude_session_leaders {
+        processes
+            .into_iter()
+            .filter(|p| p.tty.is_some() && p.pid != p.sid) // -a and -d logic
+            .collect::<Vec<_>>()
+    } else if args.terminal_processes {
+        processes
+            .into_iter()
+            .filter(|p| p.tty.is_some()) // -a logic
+            .collect::<Vec<_>>()
+    } else if args.exclude_session_leaders {
+        processes
+            .into_iter()
+            .filter(|p| p.pid != p.sid) // -d logic
+            .collect::<Vec<_>>()
+    } else {
+        processes
+    };
+
+    println!(
+        "{:<5} {:<5} {:<5} {:<5} {}",
+        "PID", "PPID", "UID", "GID", "COMMAND"
+    );
+    for proc in filtered_processes {
+        println!(
+            "{:<5} {:<5} {:<5} {:<5} {}",
+            proc.pid, proc.ppid, proc.uid, proc.gid, proc.path
+        );
     }
 }
diff --git a/sys/pslinux.rs b/sys/pslinux.rs
@@ -1,3 +1,12 @@
+//
+// Copyright (c) 2024 Hemi Labs, Inc.
+//
+// This file is part of the posixutils-rs project covered under
+// the MIT License.  For the full license text, please see the LICENSE
+// file in the root directory of this project.
+// SPDX-License-Identifier: MIT
+//
+
 use std::fs;
 use std::fs::read_to_string;
 use std::io::Error;
@@ -9,6 +18,8 @@ pub struct ProcessInfo {
     pub uid: u32,
     pub gid: u32,
     pub path: String,
+    pub tty: Option<String>, // Add TTY field for -a option
+    pub sid: i32,            // Add session ID (SID) for -d option
 }
 
 pub fn list_processes() -> Result<Vec<ProcessInfo>, Error> {
@@ -36,15 +47,19 @@ fn get_process_info(pid: i32, proc_path: &Path) -> Option<ProcessInfo> {
     let status = read_to_string(status_path).ok()?;
     let cmdline = read_to_string(cmdline_path).unwrap_or_default();
     let exe = fs::read_link(exe_path).unwrap_or_else(|_| PathBuf::from("[Permission denied]"));
-    let mut comm = String::new();
+    let comm = String::new();
 
-    if let Ok(stat) = read_to_string(stat_path) {
-        if let Some(start) = stat.find('(') {
-            if let Some(end) = stat.find(')') {
-                comm = stat[start + 1..end].to_string();
-            }
-        }
-    }
+    // Read from /proc/<pid>/stat to get the session ID and TTY number
+    let stat = read_to_string(stat_path).ok()?;
+    let stat_fields: Vec<&str> = stat.split_whitespace().collect();
+    let sid = stat_fields[5].parse().unwrap_or(0); // Session ID (SID)
+    let tty_nr = stat_fields[6].parse::<i32>().unwrap_or(0);
+
+    let tty = if tty_nr > 0 {
+        Some(format!("tty{}", tty_nr)) // Simplified TTY representation
+    } else {
+        None
+    };
 
     let mut ppid = 0;
     let mut uid = 0;
@@ -80,5 +95,7 @@ fn get_process_info(pid: i32, proc_path: &Path) -> Option<ProcessInfo> {
         uid,
         gid,
         path,
+        tty,
+        sid, // Return the session ID (SID)
     })
 }
diff --git a/sys/psmacos.rs b/sys/psmacos.rs
@@ -1,6 +1,17 @@
-use libc::{c_int, c_void, pid_t, proc_listallpids, proc_pidinfo, proc_pidpath};
+//
+// Copyright (c) 2024 Hemi Labs, Inc.
+//
+// This file is part of the posixutils-rs project covered under
+// the MIT License.  For the full license text, please see the LICENSE
+// file in the root directory of this project.
+// SPDX-License-Identifier: MIT
+//
+
+use libc::{c_int, c_void, getsid, pid_t, proc_listallpids, proc_pidinfo, proc_pidpath};
 use std::ffi::CStr;
+use std::fs;
 use std::io::Error;
+use std::os::unix::fs::MetadataExt;
 
 const PROC_PIDPATHINFO_MAXSIZE: usize = 4096;
 
@@ -10,6 +21,8 @@ pub struct ProcessInfo {
     pub uid: u32,
     pub gid: u32,
     pub path: String,
+    pub tty: Option<String>, // Add TTY field for -a option
+    pub sid: pid_t,          // Session ID (SID) for -d option
 }
 
 pub fn list_processes() -> Result<Vec<ProcessInfo>, Error> {
@@ -75,11 +88,37 @@ fn get_process_info(pid: pid_t) -> Option<ProcessInfo> {
         String::new()
     };
 
+    // Retrieve the terminal device ID (TTY)
+    let tty_dev = proc_info.e_tdev;
+
+    // Map the terminal device ID to the TTY name
+    let tty = get_tty_name(tty_dev);
+
     Some(ProcessInfo {
         pid: proc_info.pbi_pid as pid_t,
         ppid: proc_info.pbi_ppid as pid_t,
         uid: proc_info.pbi_uid,
         gid: proc_info.pbi_gid,
         path,
+        tty,                         // Add the terminal (TTY) name
+        sid: unsafe { getsid(pid) }, // Add session ID (SID)
     })
 }
+
+// Function to map a device ID to a TTY name
+fn get_tty_name(tty_dev: u32) -> Option<String> {
+    let dev_dir = "/dev/";
+    if let Ok(entries) = fs::read_dir(dev_dir) {
+        for entry in entries {
+            if let Ok(entry) = entry {
+                let path = entry.path();
+                if let Ok(metadata) = fs::metadata(&path) {
+                    if metadata.rdev() == tty_dev as u64 {
+                        return Some(path.file_name()?.to_string_lossy().into_owned());
+                    }
+                }
+            }
+        }
+    }
+    None
+}
