1
use std::{
2
    collections::{BTreeMap, BTreeSet},
3
    io::IsTerminal,
4
};
5

            
6
use anyhow::Context;
7
use clap::{Args, Parser};
8
use clap_complete::ArgValueCompleter;
9
use dialoguer::{Confirm, Editor};
10
use futures_util::SinkExt;
11
use itertools::Itertools;
12
use nix::unistd::{User, getuid};
13
use tokio_stream::StreamExt;
14

            
15
use crate::{
16
    client::commands::{erroneous_server_response, print_authorization_owner_hint},
17
    core::{
18
        completion::{mysql_database_completer, mysql_user_completer},
19
        database_privileges::{
20
            DatabasePrivilegeEdit, DatabasePrivilegeEditEntry, DatabasePrivilegeRow,
21
            DatabasePrivilegeRowDiff, DatabasePrivilegesDiff, create_or_modify_privilege_rows,
22
            diff_privileges, display_privilege_diffs, generate_editor_content_from_privilege_data,
23
            parse_privilege_data_from_editor_content, reduce_privilege_diffs,
24
        },
25
        protocol::{
26
            ClientToServerMessageStream, ListDatabasesError, ListUsersError,
27
            ModifyDatabasePrivilegesError, Request, Response,
28
            print_modify_database_privileges_output_status, request_validation::ValidationError,
29
        },
30
        types::{MySQLDatabase, MySQLUser},
31
    },
32
};
33

            
34
#[derive(Parser, Debug, Clone)]
35
pub struct EditPrivsArgs {
36
    /// The privileges to set, grant or revoke, in the format `DATABASE:USER:[+-]PRIVILEGES`
37
    ///
38
    /// This option allows for changing privileges for multiple databases and users in batch.
39
    ///
40
    /// This can not be used together with the positional `DB_NAME`, `USER_NAME` and `PRIVILEGES` arguments.
41
    #[arg(
42
      short,
43
      long,
44
      value_name = "DB_NAME:USER_NAME:[+-]PRIVILEGES",
45
      num_args = 0..,
46
      value_parser = DatabasePrivilegeEditEntry::parse_from_str,
47
      conflicts_with("single_priv"),
48
    )]
49
    pub privs: Vec<DatabasePrivilegeEditEntry>,
50

            
51
    #[command(flatten)]
52
    pub single_priv: Option<SinglePrivilegeEditArgs>,
53

            
54
    /// Print the information as JSON
55
    #[arg(short, long)]
56
    pub json: bool,
57

            
58
    /// Specify the text editor to use for editing privileges
59
    #[arg(
60
      short,
61
      long,
62
      value_name = "COMMAND",
63
      value_hint = clap::ValueHint::CommandString,
64
    )]
65
    pub editor: Option<String>,
66

            
67
    /// Disable interactive confirmation before saving changes
68
    #[arg(short, long)]
69
    pub yes: bool,
70
}
71

            
72
#[derive(Args, Debug, Clone)]
73
pub struct SinglePrivilegeEditArgs {
74
    /// The `MySQL` database to edit privileges for
75
    #[cfg_attr(not(feature = "suid-sgid-mode"), arg(add = ArgValueCompleter::new(mysql_database_completer)))]
76
    #[arg(
77
        value_name = "DB_NAME",
78
        requires = "user_name",
79
        requires = "single_priv"
80
    )]
81
    pub db_name: Option<MySQLDatabase>,
82

            
83
    /// The `MySQL` database to edit privileges for
84
    #[cfg_attr(not(feature = "suid-sgid-mode"), arg(add = ArgValueCompleter::new(mysql_user_completer)))]
85
    #[arg(value_name = "USER_NAME")]
86
    pub user_name: Option<MySQLUser>,
87

            
88
    /// The privileges to set, grant or revoke
89
    #[arg(
90
      allow_hyphen_values = true,
91
      value_name = "[+-]PRIVILEGES",
92
      value_parser = DatabasePrivilegeEdit::parse_from_str,
93
    )]
94
    pub single_priv: Option<DatabasePrivilegeEdit>,
95
}
96

            
97
async fn users_exist(
98
    server_connection: &mut ClientToServerMessageStream,
99
    privilege_diff: &BTreeSet<DatabasePrivilegesDiff>,
100
) -> anyhow::Result<BTreeMap<MySQLUser, Result<(), ListUsersError>>> {
101
    let user_list = privilege_diff
102
        .iter()
103
        .map(|diff| diff.get_user_name().clone())
104
        .collect();
105

            
106
    let message = Request::ListUsers(Some(user_list));
107
    server_connection.send(message).await?;
108

            
109
    let result = match server_connection.next().await {
110
        Some(Ok(Response::ListUsers(user_map))) => user_map,
111
        response => {
112
            erroneous_server_response(response)?;
113
            // Unreachable, but needed to satisfy the type checker
114
            BTreeMap::new()
115
        }
116
    };
117

            
118
    let result = result
119
        .into_iter()
120
        .map(|(user, user_result)| (user, user_result.map(|_| ())))
121
        .collect();
122

            
123
    Ok(result)
124
}
125

            
126
async fn databases_exist(
127
    server_connection: &mut ClientToServerMessageStream,
128
    privilege_diff: &BTreeSet<DatabasePrivilegesDiff>,
129
) -> anyhow::Result<BTreeMap<MySQLDatabase, Result<(), ListDatabasesError>>> {
130
    let database_list = privilege_diff
131
        .iter()
132
        .map(|diff| diff.get_database_name().clone())
133
        .collect();
134

            
135
    let message = Request::ListDatabases(Some(database_list));
136
    server_connection.send(message).await?;
137

            
138
    let result = match server_connection.next().await {
139
        Some(Ok(Response::ListDatabases(database_map))) => database_map,
140
        response => {
141
            erroneous_server_response(response)?;
142
            // Unreachable, but needed to satisfy the type checker
143
            BTreeMap::new()
144
        }
145
    };
146

            
147
    let result = result
148
        .into_iter()
149
        .map(|(database, db_result)| (database, db_result.map(|_| ())))
150
        .collect();
151

            
152
    Ok(result)
153
}
154

            
155
// TODO: reduce the complexity of this function
156
pub async fn edit_database_privileges(
157
    args: EditPrivsArgs,
158
    // NOTE: this is only used for backwards compat with mysql-admutils
159
    use_database: Option<MySQLDatabase>,
160
    mut server_connection: ClientToServerMessageStream,
161
) -> anyhow::Result<()> {
162
    let message = Request::ListPrivileges(use_database.clone().map(|db| vec![db]));
163

            
164
    server_connection.send(message).await?;
165

            
166
    debug_assert!(args.privs.is_empty() ^ args.single_priv.is_none());
167

            
168
    let privs = if let Some(single_priv_entry) = &args.single_priv {
169
        let database = single_priv_entry.db_name.clone().ok_or_else(|| {
170
            anyhow::anyhow!(
171
                "DB_NAME must be specified when editing privileges in single privilege mode"
172
            )
173
        })?;
174
        let user = single_priv_entry.user_name.clone().ok_or_else(|| {
175
            anyhow::anyhow!(
176
                "USER_NAME must be specified when DB_NAME is specified in single privilege mode"
177
            )
178
        })?;
179
        let privilege_edit = single_priv_entry.single_priv.clone().ok_or_else(|| {
180
            anyhow::anyhow!(
181
                "PRIVILEGES must be specified when DB_NAME is specified in single privilege mode"
182
            )
183
        })?;
184

            
185
        vec![DatabasePrivilegeEditEntry {
186
            database,
187
            user,
188
            privilege_edit,
189
        }]
190
    } else {
191
        args.privs.clone()
192
    };
193

            
194
    let existing_privilege_rows = match server_connection.next().await {
195
        Some(Ok(Response::ListPrivileges(databases))) => databases
196
            .into_iter()
197
            .filter_map(|(database_name, result)| match result {
198
                Ok(privileges) => Some(privileges),
199
                Err(err) => {
200
                    eprintln!("{}", err.to_error_message(&database_name));
201
                    eprintln!("Skipping...");
202
                    println!();
203
                    None
204
                }
205
            })
206
            .flatten()
207
            .sorted_by_key(|row| (row.db.clone(), row.user.clone()))
208
            .collect::<Vec<_>>(),
209
        Some(Ok(Response::ListAllPrivileges(privilege_rows))) => match privilege_rows {
210
            Ok(list) => list
211
                .into_iter()
212
                .sorted_by_key(|row| (row.db.clone(), row.user.clone()))
213
                .collect(),
214
            Err(err) => {
215
                server_connection.send(Request::Exit).await?;
216
                return Err(anyhow::anyhow!(err.to_error_message())
217
                    .context("Failed to list database privileges"));
218
            }
219
        },
220
        response => return erroneous_server_response(response),
221
    };
222

            
223
    let diffs: BTreeSet<DatabasePrivilegesDiff> = if privs.is_empty() {
224
        if !std::io::stdin().is_terminal() {
225
            anyhow::bail!(
226
                "Cannot launch editor in non-interactive mode. Please provide privileges via command line arguments."
227
            );
228
        }
229
        let privileges_to_change =
230
            edit_privileges_with_editor(&existing_privilege_rows, use_database.as_ref())?;
231
        diff_privileges(&existing_privilege_rows, &privileges_to_change)
232
    } else {
233
        let privileges_to_change = parse_privilege_tables(&privs)?;
234
        create_or_modify_privilege_rows(&existing_privilege_rows, &privileges_to_change)?
235
    };
236

            
237
    let database_existence_map = databases_exist(&mut server_connection, &diffs).await?;
238
    let user_existence_map = users_exist(&mut server_connection, &diffs).await?;
239

            
240
    let diffs = reduce_privilege_diffs(&existing_privilege_rows, diffs)?
241
        .into_iter()
242
        .filter(|diff| {
243
            let database_name = diff.get_database_name();
244
            let username = diff.get_user_name();
245

            
246
            if let Some(Err(err)) = database_existence_map.get(database_name) {
247
                println!("{}", err.to_error_message(database_name));
248
                println!("Skipping...");
249
                return false;
250
            }
251

            
252
            if let Some(Err(err)) = user_existence_map.get(username) {
253
                println!("{}", err.to_error_message(username));
254
                println!("Skipping...");
255
                return false;
256
            }
257

            
258
            true
259
        })
260
        .collect::<BTreeSet<_>>();
261

            
262
    if database_existence_map.values().any(|res| {
263
        matches!(
264
            res,
265
            Err(ListDatabasesError::ValidationError(
266
                ValidationError::AuthorizationError(_)
267
            ))
268
        )
269
    }) || user_existence_map.values().any(|res| {
270
        matches!(
271
            res,
272
            Err(ListUsersError::ValidationError(
273
                ValidationError::AuthorizationError(_)
274
            ))
275
        )
276
    }) {
277
        println!();
278
        print_authorization_owner_hint(&mut server_connection).await?;
279
        println!();
280
    }
281

            
282
    if diffs.is_empty() {
283
        println!("No changes to make.");
284
        server_connection.send(Request::Exit).await?;
285
        return Ok(());
286
    }
287

            
288
    println!("The following changes will be made:\n");
289
    println!("{}", display_privilege_diffs(&diffs));
290

            
291
    if std::io::stdin().is_terminal()
292
        && !args.yes
293
        && !Confirm::new()
294
            .with_prompt("Do you want to apply these changes?")
295
            .default(false)
296
            .show_default(true)
297
            .interact()?
298
    {
299
        server_connection.send(Request::Exit).await?;
300
        return Ok(());
301
    }
302

            
303
    let message = Request::ModifyPrivileges(diffs);
304
    server_connection.send(message).await?;
305

            
306
    let result = match server_connection.next().await {
307
        Some(Ok(Response::ModifyPrivileges(result))) => result,
308
        response => return erroneous_server_response(response),
309
    };
310

            
311
    print_modify_database_privileges_output_status(&result);
312

            
313
    if result.values().flatten().any(|(_, res)| {
314
        matches!(
315
            res,
316
            Err(ModifyDatabasePrivilegesError::UserValidationError(
317
                ValidationError::AuthorizationError(_)
318
            ) | ModifyDatabasePrivilegesError::DatabaseValidationError(
319
                ValidationError::AuthorizationError(_)
320
            ))
321
        )
322
    }) {
323
        print_authorization_owner_hint(&mut server_connection).await?;
324
    }
325

            
326
    server_connection.send(Request::Exit).await?;
327

            
328
    if result.values().flatten().any(|(_, res)| res.is_err()) {
329
        std::process::exit(1);
330
    }
331

            
332
    Ok(())
333
}
334

            
335
fn parse_privilege_tables(
336
    privs: &[DatabasePrivilegeEditEntry],
337
) -> anyhow::Result<BTreeSet<DatabasePrivilegeRowDiff>> {
338
    debug_assert!(!privs.is_empty());
339
    privs
340
        .iter()
341
        .map(|priv_edit_entry| {
342
            priv_edit_entry
343
                .as_database_privileges_diff()
344
                .context(format!(
345
                    "Failed parsing database privileges: `{priv_edit_entry}`"
346
                ))
347
        })
348
        .collect::<anyhow::Result<BTreeSet<DatabasePrivilegeRowDiff>>>()
349
}
350

            
351
fn edit_privileges_with_editor(
352
    privilege_data: &[DatabasePrivilegeRow],
353
    // NOTE: this is only used for backwards compat with mysql-admtools
354
    database_name: Option<&MySQLDatabase>,
355
) -> anyhow::Result<Vec<DatabasePrivilegeRow>> {
356
    let unix_user = User::from_uid(getuid())
357
        .context("Failed to look up your UNIX username")
358
        .and_then(|u| u.ok_or(anyhow::anyhow!("Failed to look up your UNIX username")))?;
359

            
360
    let editor_content =
361
        generate_editor_content_from_privilege_data(privilege_data, &unix_user.name, database_name);
362

            
363
    // TODO: handle errors better here
364
    let result = Editor::new().extension("tsv").edit(&editor_content)?;
365

            
366
    match result {
367
        None => Ok(privilege_data.to_vec()),
368
        Some(result) => parse_privilege_data_from_editor_content(&result)
369
            .context("Could not parse privilege data from editor"),
370
    }
371
}