1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
//! This crate can be used to continuously scan over `/proc` filesystem and store it in the struct `EncoDecode. This struct is serialized and is written to the `datadir`.
//! This is a wrapper over the [procfs](https://docs.rs/procfs/0.5.3/procfs/) crate, so the compatibility of this crate depends on the compatibility of the [procfs](https://docs.rs/procfs/0.5.3/procfs/) crate.
//!
//! The stored data is of type `EncoDecode` and can be read as:
//!
//! # Examples
//!
//! ```rust
//! use std::fs::File;
//! use std::io::Read;
//! use procshot_server::EncoDecode;
//! pub fn read_test_data() {
//!         let mut file = File::open("./test_data.procshot").unwrap();
//!         let mut data = Vec::new();
//!         file.read_to_end(&mut data).unwrap();
//!         let decoded: EncoDecode = bincode::deserialize(&data[..]).unwrap_or_else(|x| panic!("Error reading saved data. This was either created with an older version of procshot, or the file is corrupt. Error is {}", x));
//!         println!("Decoded test file data: {:#?}", decoded);
//! }
//! ```

extern crate procfs;
use std::collections::HashMap;
use std::thread;
use std::time::Duration;
#[macro_use]
extern crate serde_derive;
extern crate serde;
use std::fs::File;
use std::io::Write;
use std::io::{BufRead, BufReader};

// Tmp imports

extern crate clap;
extern crate hostname;
use clap::{App, Arg, SubCommand};

/// PidStatus is the struct that holds the data that we store for each process' status. In this crate, we create a
/// ` Vec<HashMap<i32, PidStatus>>` which is a mapping of pid to its status.
#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
pub struct PidStatus {
    /// Parent pid
    pub ppid: i32,
    /// Effective uid
    pub euid: i32,
    /// The complete path to cmd_long if available.
    pub cmd_long: Vec<String>,
    /// Command run by this process.
    pub name: String,
    /// The filename of the executable, in parentheses.
    ///
    /// This is visible whether or not the executable is swapped out.
    pub cmd_short: String,
    /// PID of process tracing this process (0 if not being traced).
    pub tracerpid: i32,
    /// Number of file descriptor slots currently allocated.
    pub fdsize: u32,
    /// Current state of the process.
    pub state: String,
    /// Peak virtual memory size by kB.
    pub vmpeak: Option<u64>,
    /// Virtual memory size by kB.
    pub vmsize: Option<u64>,
    /// Resident Set Size: number of pages the process has in real memory.
    ///
    /// This is just the pages which count toward text,  data,  or stack space.
    /// This does not include pages which have not been demand-loaded in, or which are swapped out.
    pub rss_pages: i64,
    /// Gets the Resident Set Size (in bytes)
    pub rss_bytes: i64,
    /// Current soft limit in bytes on the rss of the process; see the description of RLIMIT_RSS in
    /// getrlimit(2).
    pub rsslim_bytes: u64,
    /// CPU number last executed on.
    ///
    /// (since Linux 2.2.8)
    pub processor_last_executed: Option<i32>,
    // Amount of time that this process has been scheduled in user mode, measured in clock ticks
    /// (divide by [`ticks_per_second()`].
    ///
    /// This includes guest time, guest_time (time spent running a virtual CPU, see below), so that
    /// applications that are not aware of the guest time field  do not lose that time from their
    /// calculations.
    pub utime: u64,
    /// Amount of time that this process has been scheduled in kernel mode, measured in clock ticks
    /// (divide by [`ticks_per_second()`]).
    pub stime: u64,
    /// Holds the user CPU usage by that process.
    pub user_cpu_usage: f64,
    /// Holds the sys CPU usage by that process.    
    pub sys_cpu_usage: f64,
}

/// EncodDecode is the struct that we use to hold additional metadata and write to disk as
/// serialized data of the form `let enc encoded: Vec<u8> = bincode::serialize(&encodecode).unwrap();`.
#[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
pub struct EncoDecode {
    pub hostname: String,
    /// Vector of hashmap of pid to the pidstats.
    pub pid_map_list: HashMap<i32, PidStatus>,
    /// The epoch time at which the stats were recorded
    pub time_epoch: u64,
    /// Can be used for sampling
    pub delay: u64,
    /// The cumilative CPU time in jiffies.
    pub total_cpu_time: u64,
}

/// scan_proc continuously scans /proc and records all the processes.
/// scan_proc omits the pids if status.vmpeak == None || prc.stat.rss == 0 || status.pid < 0.
/// One file is created for each iteration and sleeps for `delay` seconds after each iteration.
/// The example in the description can be used as a reference to read the stored struct.
pub fn scan_proc(delay: u64, host: String, datadir: &'static str) {
    print!("Starting procshot server with delay set as {}", delay);

    let mut previous_stats: Option<HashMap<i32, PidStatus>> = None;
    let mut previous_cpu_time: u64 = 0;
    // Starts the continuous iteration over /proc
    loop {
        let mut pid_map_hash: HashMap<i32, PidStatus> = HashMap::new(); //Vec::new();
        let time_epoch = std::time::SystemTime::now()
            .duration_since(std::time::SystemTime::UNIX_EPOCH)
            .unwrap()
            .as_secs();
        let total_cpu_time = match read_proc_stat() {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Cannot read from /proc/stat, error is:: {:?}", e);
                continue;
            }
        };

        // Iterate over all processess
        for prc in procfs::all_processes() {
            let status = prc.status().unwrap_or_else(|_| dummy_pid_status());
            if status.vmpeak == None || prc.stat.rss == 0 || status.pid < 0 {
                continue;
            }
            let s = PidStatus {
                ppid: status.ppid,
                euid: status.euid,
                cmd_long: prc
                    .cmdline()
                    .unwrap_or_else(|_| vec!["No cmd_long found".to_string()]),
                name: status.name,
                cmd_short: prc.stat.comm.clone(),
                tracerpid: status.tracerpid,
                fdsize: status.fdsize,
                state: status.state,
                vmpeak: status.vmpeak,
                vmsize: status.vmsize,
                rss_pages: prc.stat.rss,
                rss_bytes: prc.stat.rss_bytes(),
                rsslim_bytes: prc.stat.rsslim,
                processor_last_executed: prc.stat.processor,
                utime: prc.stat.utime,
                stime: prc.stat.stime,
                user_cpu_usage: get_cpu_usage(
                    "user".to_string(),
                    status.pid,
                    &previous_stats,
                    prc.stat.utime,
                    total_cpu_time,
                    previous_cpu_time,
                ),
                sys_cpu_usage: get_cpu_usage(
                    "system".to_string(),
                    status.pid,
                    &previous_stats,
                    prc.stat.stime,
                    total_cpu_time,
                    previous_cpu_time,
                ),
            };

            // let mut pidmap: HashMap<i32, PidStatus> = HashMap::new();
            pid_map_hash.insert(status.pid, s);
        }
        previous_stats = Some(pid_map_hash.clone());
        previous_cpu_time = total_cpu_time;

        let encodecode: EncoDecode = EncoDecode {
            hostname: host.clone(),
            pid_map_list: pid_map_hash,
            delay: delay,
            time_epoch: time_epoch,
            total_cpu_time: total_cpu_time,
        };
        let encoded: Vec<u8> = bincode::serialize(&encodecode).unwrap();
        // println!("DECODED VALUES:: {:#?}", decoded);
        //assert_eq!(pids, decoded);
        let file = File::create(format! {"{}/{}.procshot", datadir, time_epoch});
        match file {
            Err(e) => eprintln!("Cannot create file!, err: {}", e),
            Ok(mut f) => {
                f.write_all(&encoded).unwrap();
            }
        }
        thread::sleep(Duration::from_secs(delay));
    }
}

/// get_cpu_usage calculates cpu usage for user/system.
/// user_util = 100 * (utime_after - utime_before) / (time_total_after - time_total_before);
/// sys_util = 100 * (stime_after - stime_before) / (time_total_after - time_total_before);
fn get_cpu_usage(
    type_of: String,
    pid: i32,
    previous: &Option<HashMap<i32, PidStatus>>,
    current_type_time: u64,
    current_cpu_time: u64,
    previous_cpu_time: u64,
) -> f64 {
    match type_of.as_ref() {
        "user" => match previous {
            Some(x) => match x.get(&pid) {
                Some(p) => {
                    100 as f64 * (current_type_time as f64 - p.utime as f64) / (current_cpu_time as f64 - previous_cpu_time as f64)
                }
                None => {
                    0.0
                }
            },
            None => {
                0.0
            }
        },
        "system" => match previous {
            Some(x) => match x.get(&pid) {
                Some(p) => {
                    100 as f64 * (current_type_time as f64 - p.stime as f64)
                        / (current_cpu_time as f64 - previous_cpu_time as f64)
                }
                None => 0.0,
            },
            None => 0.0,
        },
        _ => {
            println!("Keyword not supported!");
            0.0
        }
    }
}

/// Reads and parses /proc/stat's first line for calculating cpu percentage
fn read_proc_stat() -> Result<u64, std::io::Error> {
    let f = match File::open("/proc/stat") {
        Ok(somefile) => somefile,
        Err(e) => return Err(e),
    };

    let mut reader_itr = BufReader::new(f).lines();
    let first_line = match reader_itr.next() {
        // next returns an Option<Result<>> type, and hence the nested some(ok())
        Some(total_string) => match total_string {
            Ok(s) => s,
            Err(e) => return Err(e),
        },
        None => {
            return Err(std::io::Error::new(
                std::io::ErrorKind::Other,
                "Cannot read the first line from /proc/stat.",
            ))
        }
    };
    let total_vector = first_line
        .split("cpu") // Split at "cpu"
        .collect::<Vec<&str>>()[1] // Skip 0th element
        .split(" ") // Split at " "
        .filter(|&x| x != "") // filter empty lines
        .collect::<Vec<&str>>(); // collect
    let mut total: u64 = 0;
    for i in total_vector {
        total += i.parse::<u64>().unwrap();
    }
    Ok(total)
}

///dummy_status is used to return a dummy procfs::Status struct
fn dummy_pid_status() -> procfs::Status {
    let ds = "Dummy because unwrap failed".to_string();
    procfs::Status {
        name: ds.clone(),
        umask: Some(std::u32::MAX),
        state: ds.clone(),
        tgid: -1,
        ngid: Some(-1),
        pid: -1,
        ppid: -1,
        tracerpid: -1,
        ruid: -1,
        euid: -1,
        suid: -1,
        fuid: -1,
        rgid: -1,
        egid: -1,
        sgid: -1,
        fgid: -1,
        fdsize: std::u32::MAX,
        groups: vec![-1],
        nstgid: Some(vec![-1]),
        nspid: Some(vec![-1]),
        nspgid: Some(vec![-1]),
        nssid: Some(vec![-1]),
        vmpeak: Some(std::u64::MAX),
        vmsize: Some(std::u64::MAX),
        vmlck: Some(std::u64::MAX),
        vmpin: Some(std::u64::MAX),
        vmhwm: Some(std::u64::MAX),
        vmrss: Some(std::u64::MAX),
        rssanon: Some(std::u64::MAX),
        rssfile: Some(std::u64::MAX),
        rssshmem: Some(std::u64::MAX),
        vmdata: Some(std::u64::MAX),
        vmstk: Some(std::u64::MAX),
        vmexe: Some(std::u64::MAX),
        vmlib: Some(std::u64::MAX),
        vmpte: Some(std::u64::MAX),
        vmswap: Some(std::u64::MAX),
        hugetblpages: Some(std::u64::MAX),
        threads: std::u64::MAX,
        sigq: (std::u64::MAX, std::u64::MAX),
        sigpnd: std::u64::MAX,
        shdpnd: std::u64::MAX,
        sigblk: std::u64::MAX,
        sigign: std::u64::MAX,
        sigcgt: std::u64::MAX,
        capinh: std::u64::MAX,
        capprm: std::u64::MAX,
        capeff: std::u64::MAX,
        capbnd: Some(std::u64::MAX),
        capamb: Some(std::u64::MAX),
        nonewprivs: Some(std::u64::MAX),
        seccomp: Some(std::u32::MAX),
        speculation_store_bypass: Some(ds.clone()),
        cpus_allowed: Some(vec![std::u32::MAX]),
        cpus_allowed_list: Some(vec![(std::u32::MAX, std::u32::MAX)]),
        mems_allowed: Some(vec![std::u32::MAX]),
        mems_allowed_list: Some(vec![(std::u32::MAX, std::u32::MAX)]),
        voluntary_ctxt_switches: Some(std::u64::MAX),
        nonvoluntary_ctxt_switches: Some(std::u64::MAX),
    }
}

/// Config struct holds the user input when running the server. It is a bad design to hold the client's option as well in the same struct, but
/// as of now, it is here.
#[derive(Debug)]
pub struct Config {
    /// hostname of the server. This is derived by this crate from the [hostname](https://docs.rs/hostname/0.1.5/hostname/) crate.
    pub hostname: String,
    /// Delay decides how many seconds to sleep after each iteration of scanning /proc
    pub delay: u64,
    /// If true, runs as server. Defaults to false. Pass the subcommand `server` to set it to true.
    pub server: bool,
    /// The time from which the client can fetch data to process.
    pub client_time_from: String,
    /// Sort the processed data by whatever the user wants.
    pub client_sort_by: String,
}

/// Returns a new config object. This also gives the following command line argument options.
/// # Examples
/// Here are the cli options used to populate the struct.
/// > sudo target/debug/procshot --help
///
/// ```bash
/// procshot 1.0
/// nohupped@gmail.com
/// Snapshots proc periodically. All the options except delay works when 'server' option is not used.
///
/// USAGE:
///     procshot [FLAGS] [OPTIONS] [SUBCOMMAND]

/// FLAGS:
///     -h, --help       Prints help information
///     -o               Sort result by Memory or CPU. Accepted values are...
///     -t               Read stats from a specific time. Accepted format: 2015-09-05 23:56:04
///     -V, --version    Prints version information
///
/// OPTIONS:
///     -d, --delay <delay>      Sets delay in seconds before it scans /proc every time. [default: 60]
///
/// SUBCOMMANDS:
///     help      Prints this message or the help of the given subcommand(s)
///     server    Decides whether to run as server or client
impl Config {
    pub fn new() -> Self {
        let matches = App::new("procshot")
                        .version("1.0")
                        .author("nohupped@gmail.com")
                        .about("Snapshots proc periodically. All the options except delay works when 'server' option is not used.")
                        .arg(Arg::with_name("delay")
                            .short("d")
                            .long("delay")
                            .default_value("60")
                            .help("Sets delay in seconds before it scans /proc every time."))
                        .subcommand(SubCommand::with_name("server")
                            .about("Runs as server and records stats."))
                        .arg(Arg::with_name("time_from")
                            .short("t")
                            .help("Read stats from a specific time. Accepted format: 2015-09-05 23:56:04")
                            )
                        .arg(Arg::with_name("order_by")
                            .short("o")
                            .help("Sort result by Memory or CPU. Accepted values are...") // Todo here
                            )
                        .get_matches();

        Config {
            hostname: hostname::get_hostname().unwrap().to_string(),
            delay: matches
                .value_of("delay")
                .unwrap_or("60")
                .parse()
                .unwrap_or(60),
            server: match matches.subcommand_matches("server") {
                Some(_) => true,
                None => false,
            },
            client_time_from: matches.value_of("time_from").unwrap_or("").to_string(),
            client_sort_by: matches.value_of("order_by").unwrap_or("m").to_string(),
        }
    }
}
/// Checks if the program is run as sudo (root) user. This doesn't check if the user has the privilege to read over all of /proc or write to the datadir
/// but just checks if the uid passed to this is 0, and returns a `Result`
///
/// # Examples
///
///```rust
///
/// use procshot_server::check_sudo;
/// use std::process;
///
/// fn main() {
///     match check_sudo(0) { // Can also use get_current_uid() from the `users` crate
///         Err(e) => {
///             eprintln!("Error encountered checking privileges, {:?}", e);
///             process::exit(1);
///         },
///         _ => (),
///     }
/// }
///```
pub fn check_sudo(uid: u32) -> Result<(), &'static str> {
    match uid == 0 {
        true => Ok(()),
        false => Err("Error: Run as root."),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_check_sudo_privileged() {
        match check_sudo(0) {
            Ok(()) => (),
            Err(e) => panic!("Test failed, {:?}", e),
        }
    }

    #[test]
    #[should_panic]
    fn test_check_sudo_non_privileged() {
        match check_sudo(10) {
            Ok(()) => (),
            Err(e) => panic!("Test failed, {:?}", e),
        }
    }
}