use std::{convert::TryFrom, env, io, process::exit, sync::Arc}; use clap::{App as Argparse, AppSettings as ArgParseSettings, Arg, ArgMatches, SubCommand}; use futures::{executor::block_on, StreamExt}; use rustyline::{ completion::{Completer, Pair}, error::ReadlineError, highlight::{Highlighter, MatchingBracketHighlighter}, hint::{Hinter, HistoryHinter}, validate::{MatchingBracketValidator, Validator}, CompletionType, Config, Context, EditMode, Editor, OutputStreamType, }; use rustyline_derive::Helper; use syntect::{ easy::HighlightLines, highlighting::{Style, ThemeSet}, parsing::SyntaxSet, util::{as_24_bit_terminal_escaped, LinesWithEndings}, }; use matrix_sdk_base::{ events::EventType, identifiers::{RoomId, UserId}, RoomInfo, Store, }; #[derive(Clone)] struct Inspector { store: Store, printer: Printer, } #[derive(Helper)] struct InspectorHelper { store: Store, _highlighter: MatchingBracketHighlighter, _validator: MatchingBracketValidator, _hinter: HistoryHinter, } impl InspectorHelper { const EVENT_TYPES: &'static [&'static str] = &[ "m.room.aliases", "m.room.avatar", "m.room.canonical_alias", "m.room.create", "m.room.encryption", "m.room.guest_access", "m.room.history_visibility", "m.room.join_rules", "m.room.name", "m.room.power_levels", "m.room.tombstone", "m.room.topic", ]; fn new(store: Store) -> Self { Self { store, _highlighter: MatchingBracketHighlighter::new(), _validator: MatchingBracketValidator::new(), _hinter: HistoryHinter {}, } } fn complete_event_types(&self, arg: Option<&&str>) -> Vec { Self::EVENT_TYPES .iter() .map(|t| Pair { display: t.to_string(), replacement: format!("{} ", t), }) .filter(|r| { if let Some(arg) = arg { r.replacement.starts_with(arg) } else { true } }) .collect() } fn complete_rooms(&self, arg: Option<&&str>) -> Vec { let rooms: Vec = block_on(async { self.store.get_room_infos().await.collect().await }); rooms .into_iter() .map(|r| Pair { display: r.room_id.to_string(), replacement: format!("{} ", r.room_id.to_string()), }) .filter(|r| { if let Some(arg) = arg { r.replacement.starts_with(arg) } else { true } }) .collect() } } impl Completer for InspectorHelper { type Candidate = Pair; fn complete( &self, line: &str, pos: usize, _: &Context<'_>, ) -> Result<(usize, Vec), ReadlineError> { let args: Vec<&str> = line.split_ascii_whitespace().collect(); let commands = vec![ ("get-state", "get a state event in the given room"), ( "get-profiles", "get all the stored profiles in the given room", ), ("list-rooms", "list all rooms"), ( "get-members", "get all the membership events in the given room", ), ] .iter() .map(|(r, d)| Pair { display: format!("{} ({})", r, d), replacement: format!("{} ", r.to_string()), }) .collect(); if args.is_empty() { Ok((pos, commands)) } else if args.len() == 1 { if (args[0] == "get-state" || args[0] == "get-members" || args[0] == "get-profiles") && line.ends_with(' ') { Ok((args[0].len() + 1, self.complete_rooms(args.get(1)))) } else { Ok(( 0, commands .into_iter() .filter(|c| c.replacement.starts_with(args[0])) .collect(), )) } } else if args.len() == 2 { if args[0] == "get-state" { if line.ends_with(' ') { Ok(( args[0].len() + args[1].len() + 2, self.complete_event_types(args.get(2)), )) } else { Ok((args[0].len() + 1, self.complete_rooms(args.get(1)))) } } else if args[0] == "get-members" || args[0] == "get-profiles" { Ok((args[0].len() + 1, self.complete_rooms(args.get(1)))) } else { Ok((pos, vec![])) } } else if args.len() == 3 { if args[0] == "get-state" { Ok(( args[0].len() + args[1].len() + 2, self.complete_event_types(args.get(2)), )) } else { Ok((pos, vec![])) } } else { Ok((pos, vec![])) } } } impl Hinter for InspectorHelper { type Hint = String; } impl Highlighter for InspectorHelper {} impl Validator for InspectorHelper {} #[derive(Clone, Debug)] struct Printer { ps: Arc, ts: Arc, } impl Printer { fn new() -> Self { Self { ps: SyntaxSet::load_defaults_newlines().into(), ts: ThemeSet::load_from_folder("/home/poljar/local/cfg/bat/themes") .expect("Couldn't load themes") .into(), } } fn pretty_print_struct(&self, data: &impl std::fmt::Debug) { let data = format!("{:#?}", data); let syntax = self.ps.find_syntax_by_extension("rs").unwrap(); let mut h = HighlightLines::new(syntax, &self.ts.themes["Forest Night"]); for line in LinesWithEndings::from(&data) { let ranges: Vec<(Style, &str)> = h.highlight(line, &self.ps); let escaped = as_24_bit_terminal_escaped(&ranges[..], false); print!("{}", escaped); } // Clear the formatting println!("\x1b[0m"); } } impl Inspector { fn new(database_path: &str) -> Self { let printer = Printer::new(); let store = Store::open_default(database_path); Self { store, printer } } async fn run(&self, matches: ArgMatches<'_>) { match matches.subcommand() { ("get-profiles", args) => { let args = args.expect("No args provided for get-state"); let room_id = RoomId::try_from(args.value_of("room-id").unwrap()).unwrap(); self.get_profiles(room_id).await; } ("get-members", args) => { let args = args.expect("No args provided for get-state"); let room_id = RoomId::try_from(args.value_of("room-id").unwrap()).unwrap(); self.get_members(room_id).await; } ("list-rooms", _) => self.list_rooms().await, ("get-state", args) => { let args = args.expect("No args provided for get-state"); let room_id = RoomId::try_from(args.value_of("room-id").unwrap()).unwrap(); let event_type = EventType::try_from(args.value_of("event-type").unwrap()).unwrap(); self.get_state(room_id, event_type).await; } _ => unreachable!(), } } async fn list_rooms(&self) { let rooms: Vec = self.store.get_room_infos().await.collect().await; self.printer.pretty_print_struct(&rooms); } async fn get_profiles(&self, room_id: RoomId) { let joined: Vec = self .store .get_joined_user_ids(&room_id) .await .collect() .await; for member in joined { let event = self.store.get_profile(&room_id, &member).await; self.printer.pretty_print_struct(&event); } } async fn get_members(&self, room_id: RoomId) { let joined: Vec = self .store .get_joined_user_ids(&room_id) .await .collect() .await; for member in joined { let event = self.store.get_member_event(&room_id, &member).await; self.printer.pretty_print_struct(&event); } } async fn get_state(&self, room_id: RoomId, event_type: EventType) { self.printer .pretty_print_struct(&self.store.get_state_event(&room_id, event_type, "").await); } async fn parse_and_run(&self, input: &str) { let argparse = Argparse::new("state-inspector") .global_setting(ArgParseSettings::DisableHelpFlags) .global_setting(ArgParseSettings::DisableVersion) .global_setting(ArgParseSettings::VersionlessSubcommands) .global_setting(ArgParseSettings::NoBinaryName) .setting(ArgParseSettings::SubcommandRequiredElseHelp) .subcommand(SubCommand::with_name("list-rooms")) .subcommand(SubCommand::with_name("get-members").arg( Arg::with_name("room-id").required(true).validator(|r| { RoomId::try_from(r) .map(|_| ()) .map_err(|_| "Invalid room id given".to_owned()) }), )) .subcommand(SubCommand::with_name("get-profiles").arg( Arg::with_name("room-id").required(true).validator(|r| { RoomId::try_from(r) .map(|_| ()) .map_err(|_| "Invalid room id given".to_owned()) }), )) .subcommand( SubCommand::with_name("get-state") .arg(Arg::with_name("room-id").required(true).validator(|r| { RoomId::try_from(r) .map(|_| ()) .map_err(|_| "Invalid room id given".to_owned()) })) .arg(Arg::with_name("event-type").required(true).validator(|e| { EventType::try_from(e) .map(|_| ()) .map_err(|_| "Invalid event type".to_string()) })), ); match argparse.get_matches_from_safe(input.split_ascii_whitespace()) { Ok(m) => { self.run(m).await; } Err(e) => { println!("{}", e); } } } } fn main() -> io::Result<()> { let database_path = match env::args().nth(1) { Some(a) => a, _ => { eprintln!("Usage: {} ", env::args().next().unwrap()); exit(1) } }; let inspector = Inspector::new(&database_path); let config = Config::builder() .history_ignore_space(true) .completion_type(CompletionType::List) .edit_mode(EditMode::Emacs) .output_stream(OutputStreamType::Stdout) .build(); let helper = InspectorHelper::new(inspector.store.clone()); let mut rl = Editor::::with_config(config); rl.set_helper(Some(helper)); while let Ok(input) = rl.readline(">> ") { rl.add_history_entry(input.as_str()); block_on(inspector.parse_and_run(input.as_str())); } Ok(()) }