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

// Copyright 2021 - 2021, Rupansh Sekar and the Kilogramme (TBD) contributors
// SPDX-License-Identifier: MPL-2.0
mod config;
mod consts;
/// High level Error related wrapper Enums
mod errors;
mod helpers;
/// Command Handlers
mod plugins;
mod userbot;

use config::UserBotConfig;

use grammers_client::Update;
use helpers::time;
use simplelog::*;
use std::fs::File;
use tokio::sync::watch;
use userbot::*;

// TODO: Documentation, Kang
fn main() {
    let config = UserBotConfig::from_file();
    if config.is_err() {
        eprintln!("invalid config!! {}", config.err().unwrap());
        return;
    }
    let config = config.unwrap();

    // Initialize Logging
    let mut loggers: Vec<Box<dyn SharedLogger>> = vec![TermLogger::new(
        LevelFilter::Info,
        Config::default(),
        TerminalMode::Mixed,
        ColorChoice::Auto,
    )];
    if config.options.file_log {
        loggers.push(WriteLogger::new(
            LevelFilter::Debug,
            Config::default(),
            File::create(consts::LOG_FILE).unwrap(),
        ))
    }

    simplelog::CombinedLogger::init(loggers).unwrap();

    log::info!("started logging @ {}", time::epoch_ms());

    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(async_main(config));

    log::info!("bot stopped @{}", time::epoch_ms());
}

async fn async_main(config: UserBotConfig) {
    // Initialize userbot
    let store = UserBotStore::from_config(config).await;
    if let Err(e) = store {
        log::error!("failed to initialize bot {}", e);
        return;
    }
    let mut store = store.unwrap();

    let (tx, mut rx) = watch::channel::<()>(());

    let ctrlc_res = ctrlc::set_handler(move || {
        log::info!("shutting down bot...");
        tx.send(()).unwrap();
    });

    if ctrlc_res.is_err() {
        log::warn!("failed to set ctrl+c handler. Bot won't save session before quitting!");
    }

    // Start fetching updates
    while let Some(update) = tokio::select! {
        biased;
        res = rx.changed() => {
            res.unwrap();
            return;
        },
        update = store.next_update() => {
            update
        }
    } {
        let bot = UserBot::from_store(&store);
        if let Update::NewMessage(message) = update {
            let mut bot = bot.clone();
            tokio::task::spawn(async move {
                if let Err(e) = bot.update_handler(message).await {
                    log::warn!("error handling update: {}", e);
                }
            });
        }
    }
}