1
use anyhow::Context;
2
use chrono::{Duration, Utc};
3
use clap::{CommandFactory, Parser};
4
use clap_complete::{Shell, generate};
5

            
6
use roowho2_lib::{
7
    proto::WhodStatusUpdate,
8
    server::varlink_api::{
9
        DEFAULT_CLIENT_SERVER_SOCKET_PATH, VarlinkRwhodClientError, VarlinkRwhodClientProxy,
10
    },
11
    version,
12
};
13

            
14
/// Show host status of local machines.
15
///
16
/// `ruptime` gives a status line like uptime for each machine on the local network;
17
/// these are formed from packets broadcast by each host on the network once a minute.
18
///
19
/// Machines for which no status report has been received for a while (11 minutes by default,
20
/// see `--uptime-timeout`) are shown as being down.
21
#[derive(Debug, Parser)]
22
#[command(
23
  author = "Programvareverkstedet <projects@pvv.ntnu.no>",
24
  version,
25
  long_version = version::LONG_VERSION
26
)]
27
pub struct Args {
28
    /// Users idle an hour or more are not counted unless the `-a` flag is given.
29
    #[arg(long, short, conflicts_with = "idle_timeout")]
30
    all: bool,
31

            
32
    /// Only count users idle for less than this many minutes.
33
    /// Cannot be used together with `--all`.
34
    #[arg(long, value_name = "MINUTES", default_value_t = DEFAULT_IDLE_TIMEOUT_MINUTES)]
35
    idle_timeout: u32,
36

            
37
    /// Consider a machine down if no status report has been received from it for this many minutes.
38
    #[arg(long, value_name = "MINUTES", default_value_t = DEFAULT_UPTIME_TIMEOUT_MINUTES)]
39
    uptime_timeout: u32,
40

            
41
    /// Sort by load average.
42
    #[arg(long, short, group = "sort_mode")]
43
    load: bool,
44

            
45
    /// Reverses the sort order.
46
    #[arg(long, short)]
47
    reverse: bool,
48

            
49
    /// Sort by uptime.
50
    #[arg(long, short, group = "sort_mode")]
51
    time: bool,
52

            
53
    /// Sort by number of users.
54
    #[arg(long, short, group = "sort_mode")]
55
    users: bool,
56

            
57
    /// Print the output with the old formatting
58
    #[arg(long, short)]
59
    old: bool,
60

            
61
    /// Output in JSON format
62
    #[arg(long, short, conflicts_with_all = ["sort_mode", "reverse"])]
63
    json: bool,
64

            
65
    /// Generate shell completion scripts for the specified shell
66
    /// and print them to stdout.
67
    #[arg(long, value_enum, hide = true)]
68
    completions: Option<Shell>,
69
}
70

            
71
/// Users idle for longer than this are not counted by default.
72
const DEFAULT_IDLE_TIMEOUT_MINUTES: u32 = 60;
73

            
74
/// Machines that haven't reported in for longer than this are shown as down by default.
75
const DEFAULT_UPTIME_TIMEOUT_MINUTES: u32 = 11;
76

            
77
#[tokio::main]
78
async fn main() -> anyhow::Result<()> {
79
    let args = Args::parse();
80

            
81
    if let Some(shell) = args.completions {
82
        generate(
83
            shell,
84
            &mut Args::command(),
85
            "ruptime",
86
            &mut std::io::stdout(),
87
        );
88
        return Ok(());
89
    }
90

            
91
    let max_idle_seconds = if args.all {
92
        None
93
    } else {
94
        Some(i64::from(args.idle_timeout) * 60)
95
    };
96

            
97
    let mut conn = zlink::tokio::unix::connect(DEFAULT_CLIENT_SERVER_SOCKET_PATH)
98
        .await
99
        .expect("Failed to connect to rwhod server");
100

            
101
    let mut reply = conn
102
        .ruptime(max_idle_seconds)
103
        .await
104
        .context("Failed to send rwho request")?
105
        .map_err(|e| match e {
106
            VarlinkRwhodClientError::Disabled => {
107
                anyhow::anyhow!("The rwhod service is disabled on the server")
108
            }
109
            VarlinkRwhodClientError::TimedOut => {
110
                anyhow::anyhow!("The rwhod service timed out while processing the request")
111
            }
112
            VarlinkRwhodClientError::InvalidRequest => {
113
                anyhow::anyhow!("The rwhod service could not process the request, please check the logs or report the error to your system administrators")
114
            }
115
        })?;
116

            
117
    sort_entries(&mut reply, args.load, args.time, args.users, args.reverse);
118

            
119
    if args.json {
120
        println!("{}", serde_json::to_string_pretty(&reply).unwrap());
121
    // } else if args.old {
122
    //     for entry in &reply {
123
    //         let line = old_format_machine_entry(args.all, entry);
124
    //         println!("{}", line);
125
    //     }
126
    } else {
127
        for entry in &reply {
128
            let line = old_format_machine_entry(entry, args.uptime_timeout);
129
            println!("{}", line);
130
        }
131
    }
132

            
133
    Ok(())
134
}
135

            
136
fn sort_entries(
137
    entries: &mut [WhodStatusUpdate],
138
    sort_by_load: bool,
139
    sort_by_time: bool,
140
    sort_by_users: bool,
141
    reverse: bool,
142
) {
143
    entries.sort_by(|entry1, entry2| {
144
        let ordering = if sort_by_load {
145
            let load1 = entry1.load_average.0 + entry1.load_average.1 + entry1.load_average.2;
146
            let load2 = entry2.load_average.0 + entry2.load_average.1 + entry2.load_average.2;
147
            load1
148
                .partial_cmp(&load2)
149
                .unwrap_or(std::cmp::Ordering::Equal)
150
        } else if sort_by_time {
151
            let uptime1 = Utc::now() - entry1.sendtime;
152
            let uptime2 = Utc::now() - entry2.sendtime;
153
            uptime1.cmp(&uptime2)
154
        } else if sort_by_users {
155
            let users1 = entry1.users.len();
156
            let users2 = entry2.users.len();
157
            users1.cmp(&users2)
158
        } else {
159
            entry1.hostname.cmp(&entry2.hostname)
160
        };
161

            
162
        if reverse {
163
            ordering.reverse()
164
        } else {
165
            ordering
166
        }
167
    });
168
}
169

            
170
fn old_format_machine_entry(entry: &WhodStatusUpdate, uptime_timeout_minutes: u32) -> String {
171
    let time_since_last_ping = Utc::now() - entry.sendtime;
172
    let is_up = time_since_last_ping <= Duration::minutes(i64::from(uptime_timeout_minutes));
173

            
174
    let uptime = Utc::now() - entry.boot_time;
175
    let days = uptime.num_days();
176
    let hours = uptime.num_hours() % 24;
177
    let minutes = uptime.num_minutes() % 60;
178

            
179
    let uptime_str = if days > 0 {
180
        format!("{:3}+{:02}:{:02}", days, hours, minutes)
181
    } else if uptime.num_seconds() < 0 || days > 999 {
182
        "    ??:??".to_string()
183
    } else {
184
        format!("    {:2}:{:02}", hours, minutes)
185
    };
186

            
187
    let user_count = entry.users.len();
188

            
189
    format!(
190
        "{:<12.12} {} {},  {:4} user{}  load {:>4.2}, {:>4.2}, {:>4.2}",
191
        entry.hostname,
192
        if is_up { "up" } else { "down" },
193
        uptime_str,
194
        user_count,
195
        if user_count == 1 { ", " } else { "s," },
196
        entry.load_average.0 as f32 / 100.0,
197
        entry.load_average.1 as f32 / 100.0,
198
        entry.load_average.2 as f32 / 100.0,
199
    )
200
}