pluralsync/src/main.rs

use std::fs;
#[cfg(feature = "avatar")]
use std::fs::remove_file;
use std::fs::create_dir;
use std::path::Path;
use std::collections::HashMap;
#[cfg(feature = "avatar")]
use std::process::Command;

use reqwest::RequestBuilder;
use reqwest::header::{USER_AGENT, AUTHORIZATION};
#[cfg(feature = "fedi")]
use reqwest::multipart::{Part, Form};

use serde::Serialize;
use serde_json::Value;

use dirs;

mod configdata;
use configdata::*;

mod systemdata;
use systemdata::*;

mod clap_ps;
use clap_ps::*;
use clap::Parser;

const PK_URL: &str = "https://api.pluralkit.me/v2";
const SP_URL: &str = "https://api.apparyllis.com/v1";

const EXAMPLE_JSON: &str = r#"{
  "pk_key": "// Pluralkit token",
  "sp_key": "// Simplplural token",
  "avatar_module": {
    "enabled": false,
    "avatar_folder": "// Folder to grab profile pictures, they follow they member1member2...memberX.png format",
    "avatar_output_path": "// Path for the copied selected avatar"
  },
  "disc_module": {
    "enabled": false,
    "token": "// Discord user token",
    "python_path": "// Path to the python executable",
    "script_path": "// Path to updatediscordavatar.py"
  },
  "fedi_module": {
    "enabled": false,
    "instance" : "// Fedi instance url",
    "token": "// Fedi bearer token"
  }
}"#;

fn main() {

    let cli = Args::parse();

    let config_path = match dirs::config_dir() {
        Some(d) => format!("{}/pluralsync", d.display()),
        None => {
            println!("Could not fetch the config directory");
            return ();
        }
    };

    let mut res: Result<(), &str> = Ok(());

    match cli.cmd {
        Commands::Sync => {
            res = sync(config_path);
        },

        // SET MEMBER
        #[cfg(all(feature = "discord", feature = "fedi"))]
        Commands::Set { members, discord, fedi } => {
            res = set_member(config_path.clone(), members);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), discord, fedi);
        },
        #[cfg(all(feature = "discord", not(feature = "fedi")))]
        Commands::Set { members, discord } => {
            res = set_member(config_path.clone(), members);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), discord, false);
        },
        #[cfg(all(not(feature = "discord"), feature = "fedi"))]
        Commands::Set { members, fedi} => {
            res = set_member(config_path.clone(), members);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), false, fedi);
        },
        #[cfg(all(not(feature = "discord"), not(feature = "fedi")))]
        Commands::Set { members } => {
            res = set_member(config_path.clone(), members);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), false, false);
        },

        // ADD MEMBER
        #[cfg(all(feature = "discord", feature = "fedi"))]
        Commands::Add { members, discord, fedi } => {
            res = add_member(config_path.clone(), members);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), discord, fedi);
        },
        #[cfg(all(feature = "discord", not(feature = "fedi")))]
        Commands::Add { members, discord } => {
            res = add_member(config_path.clone(), members);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), discord, false);
        },
        #[cfg(all(not(feature = "discord"), feature = "fedi"))]
        Commands::Add { members, fedi} => {
            res = add_member(config_path.clone(), members);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), false, fedi);
        },
        #[cfg(all(not(feature = "discord"), not(feature = "fedi")))]
        Commands::Add { members } => {
            res = add_member(config_path.clone(), members);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), false, false);
        },

        // Get MEMBER
        #[cfg(all(feature = "discord", feature = "fedi"))]
        Commands::Get { force_from, discord, fedi } => {
            let ff = match force_from {
                Some(x) => x,
                None => ForceFrom::None,
            };

            let _ = get(config_path.clone(), ff);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), discord, fedi);
        },
        #[cfg(all(feature = "discord", not(feature = "fedi")))]
        Commands::Get { force_from, discord } => {
            let ff = match force_from {
                Some(x) => x,
                None => ForceFrom::None,
            };

            let _ = get(config_path.clone(), ff);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), discord, false);
        },
        #[cfg(all(not(feature = "discord"), feature = "fedi"))]
        Commands::Get { force_from, fedi} => {
            let ff = match force_from {
                Some(x) => x,
                None => ForceFrom::None,
            };

            let _ = get(config_path.clone(), ff);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), false, fedi);
        },
        #[cfg(all(not(feature = "discord"), not(feature = "fedi")))]
        Commands::Get { force_from } => {
            let ff = match force_from {
                Some(x) => x,
                None => ForceFrom::None,
            };

            let _ = get(config_path.clone(), ff);

            #[cfg(feature = "avatar")]
            let _ = update_avatars(config_path.clone(), false, false);
        },

        Commands::Members => {
            res = memberlist(config_path);
        },
    }

    match res {
        Ok(_) => (),
        Err(e) => println!("{}", e),
    }
}

fn sync(config_path: String) -> Result<(), &'static str> {
    // Get config
    let config: Config = match get_config(&config_path) {
        Ok(c) => c,
        Err(e) => return Err(e)
    };

    // Get Pluralkit system id
    let pk_sys = pk_get_system(&config.pk_key);
    let pk_sysid = pk_sys["id"].as_str().unwrap();

    // Get Simplyplural user id
    let sp_user_id = sp_get_userid(&config.sp_key);

    // Get Simplyplural member ids
    let sp_member_ids = sp_get_memberids(&config.sp_key, &sp_user_id);

    // get members
    let pk_members = pk_get_members(&config.pk_key, pk_sysid);
    let mut members: Vec<Member> = Vec::new();
    for member in pk_members {
        let mut m = Member {
            pk_id: member["id"].as_str().unwrap().to_string(),
            sp_id: String::new(),
            name: member["name"].as_str().unwrap().to_string(),
            alias: String::new()
        };

        if member["display_name"].as_str() != None {
            m.alias = member["display_name"].as_str().unwrap().to_string();
        } else {
            m.alias = String::from(&m.name);
        }

        m.sp_id = get_sp_id(&m, &sp_member_ids);

        members.push(m);
    }

    let sys = System {
        pk_userid: pk_sysid.to_string(),
        sp_userid: sp_user_id,
        members: members.clone(),
    };

    let json = serde_json::to_string(&sys);
    let _ = fs::write(format!("{}/system.json", config_path), &json.unwrap());

    Ok(())
}

fn set_member(config_path: String, tf_members: Vec<String>) -> Result<(), &'static str> {
    let config: Config = match get_config(&config_path) {
        Ok(c) => c,
        Err(e) => return Err(e)
    };

    let system: System = get_system(&config_path);

    let mut to_front: Vec<Member> = Vec::new();
    for member in &tf_members {
        for mem in &system.members {
            if mem.name.to_lowercase() == member.to_lowercase() || mem.alias.to_lowercase() == member.to_lowercase() {
                println!("Member {member} found");
                to_front.push(mem.clone());

                break;
            }
        }
    }
    if to_front.len() == tf_members.len() {
        let fronters = get_fronters(&config.pk_key, &config.sp_key, &system, ForceFrom::None);

        pk_set_fronters(&config.pk_key, &system, &to_front, &fronters);
        sp_set_fronters(&config.sp_key, &to_front, &fronters);
    } else {
        println!("One or more members were not found. Known members:\n--------------------------");
        let _ = memberlist(config_path);
        println!("--------------------------\nIf a member is missing from the system try running \"pluralsync sync\" to refresh the local database");
    }
    Ok(())
}

fn add_member(config_path: String, tf_members: Vec<String>) -> Result<(), &'static str> {
    let config: Config = match get_config(&config_path) {
        Ok(c) => c,
        Err(e) => return Err(e)
    };

    let system: System = get_system(&config_path);

    let mut to_front: Vec<Member> = Vec::new();
    for member in &tf_members {
        for mem in &system.members {
            if mem.name.to_lowercase() == member.to_lowercase() || mem.alias.to_lowercase() == member.to_lowercase() {
                to_front.push(mem.clone());
                break;
            }
        }
    }
    if to_front.len() == tf_members.len() {
        let fronters = get_fronters(&config.pk_key, &config.sp_key, &system, ForceFrom::None);

        let mut aux: Vec<Member> = Vec::new();
        aux.append(&mut fronters.pk.clone());
        aux.append(&mut to_front);
        to_front = aux;

        pk_set_fronters(&config.pk_key, &system, &to_front, &fronters);
        sp_set_fronters(&config.sp_key, &to_front, &fronters);

        let _ = get(config_path, ForceFrom::None);
    } else {
        println!("One or more members were not found. Known members:\n--------------------------");
        let _ = memberlist(config_path);
        println!("--------------------------\nIf a member is missing from the system try running \"pluralsync sync\" to refresh the local database");
    }
    Ok(())
}

fn memberlist(config_path: String) -> Result<(), &'static str> {
    let sys = get_system(&config_path);

    for mem in sys.members {
        if mem.name != mem.alias {
            println!("{} / {}", mem.name, mem.alias);
        } else {
            println!("{}", mem.name);
        }
    }

    Ok(())

}

fn get(config_path: String, ff: ForceFrom) -> Result<Vec<String>, &'static str> {
    let config: Config = match get_config(&config_path) {
        Ok(c) => c,
        Err(e) => return Err(e)
    };
    let sys = get_system(&config_path);

    let f = get_fronters(&config.pk_key, &config.sp_key, &sys, ff);
    let mut names = Vec::new();
    for m in &f.pk {
        names.push(String::from(&m.name));
    }
    let fronters = names.join(" || ");
    println!("Currently fronting: {}", fronters);
    let _ = fs::write(format!("{}/.front", config_path), fronters);

    Ok(names)
}

#[cfg(feature = "avatar")]
#[allow(unused_variables)]
fn update_avatars(config_path: String, discord: bool, fedi: bool) -> Result<(), &'static str>{
    #[allow(unused_mut)]
    let mut config: Config = match get_config(&config_path) {
        Ok(c) => c,
        Err(e) => return Err(e)
    };

    let names = get(config_path, ForceFrom::None).unwrap();

    #[cfg(feature = "avatar")]
    avatar_module(&config, &names);

    #[cfg(feature = "discord")] {
        config.disc_module.enabled |= discord;
        discord_module(&config);
    }
    #[cfg(feature = "fedi")] {
        config.fedi_module.enabled |= fedi;
        fedi_module(&config);
    }

    Ok(())
}

#[cfg(feature = "avatar")]
fn avatar_module(config: &Config, names: &Vec<String>) {
    if config.avatar_module.enabled {
        let avatarnames = names.join("").to_lowercase() + ".png";

        if Path::new(&config.avatar_module.avatar_output_path).exists() {
            let _ = remove_file(&config.avatar_module.avatar_output_path);
        }

        if cfg!(target_os = "windows") {
            let cmdaug = format!("copy {}\\{} {}", &config.avatar_module.avatar_folder, avatarnames, &config.avatar_module.avatar_output_path);
            Command::new("cmd").args(["/C", &cmdaug]).output().expect("Avatar module error");
        } else {
            Command::new("sh").arg("-c").arg(format!("cp {}/{} {}", &config.avatar_module.avatar_folder, avatarnames, &config.avatar_module.avatar_output_path)).output().expect("Avatar module error");
        }
    }

}

#[cfg(feature = "discord")]
fn discord_module(config: &Config) {
    if config.disc_module.enabled {
        if cfg!(target_os = "windows") {
            let mut c = Command::new("cmd").args(["/C", format!("{} {}", &config.disc_module.python_path, &config.disc_module.script_path).as_str()]).spawn().expect("Discord module error");
            let _ = c.wait().expect("Error");
        } else {
            let mut c = Command::new("sh").arg("-c").arg(format!("{} {}", &config.disc_module.python_path, &config.disc_module.script_path)).spawn().expect("Discord module error");
            let _ = c.wait().expect("Error");
        }
    }
}

#[cfg(feature = "fedi")]
fn fedi_module(config: &Config) {
    if config.fedi_module.enabled {

        let client = reqwest::Client::new();
        let form = Form::new().part("avatar", Part::bytes(fs::read(config.avatar_module.avatar_output_path.clone()).unwrap()).file_name("face.png").mime_str("image/png").unwrap());
        let rb = client
            .patch(config.fedi_module.instance.clone() + "/api/v1/accounts/update_credentials")
            .multipart(form)
            .header(USER_AGENT, "Pluralsync")
            .header(AUTHORIZATION, format!("Bearer {}", &config.fedi_module.token).as_str());

        match http_request(rb) {
            Ok(_) => (),
            Err(e) => println!("{}", e.to_string()),
        }

    }
}

fn pk_get_system(key: &str) -> Value {
    let url = format!("{}/systems/@me", PK_URL);

    let res = http_get(url,key);
    return serde_json::from_str(&res.unwrap()).unwrap();
}

fn pk_get_members(key: &str, sysid: &str) -> Vec<Value> {
    let url = format!("{}/systems/{}/members", PK_URL, sysid);

    let res = http_get(url,key);
    let datas: Vec<Value> = serde_json::from_str(&res.unwrap()).unwrap();

    return datas;
}

fn pk_get_fronters(key: &str, sys: &System) -> Vec<Member> {
    let url = format!("{}/systems/{}/fronters", PK_URL, sys.pk_userid);

    let res = http_get(url,key);
    let data: Value = serde_json::from_str(&res.unwrap()).unwrap();
    let memberdata = &data["members"].as_array();

    let mut members: Vec<Member> = Vec::new();
    for member in memberdata {
        for m in member.into_iter() {
            for dbmem in &sys.members {
                if m["name"].as_str().unwrap() == dbmem.name {
                    members.push(dbmem.clone());
                }
            }
        }
    }

    return members;
}

fn pk_set_fronters(key: &str, sys: &System, to_front: &Vec<Member>, fronters: &Fronters) {
    let url = format!("{}/systems/{}/switches", PK_URL, sys.pk_userid);

    if to_front != &fronters.pk {
        let mut frontcodes = Vec::new();
        for tf in to_front {
            frontcodes.push(String::from(&tf.pk_id));
        }

        let mut body: HashMap<&str, Vec<String>> = HashMap::new();
        body.insert("members", frontcodes);

        let client = reqwest::Client::new();
        let rb = client
            .post(url)
            .json(&body)
            .header("content-type", "application/json; charset=utf-8")
            .header(USER_AGENT, "Pluralsync")
            .header(AUTHORIZATION, key);

        match http_request(rb) {
            Ok(_) => (),
            Err(e) => println!("{}", e.to_string()),
        }

    } else {
        println!("Members already fonting");
    }

}

fn sp_get_userid(key: &str) -> String {
    let url = format!("{}/me", SP_URL);

    let res = http_get(url,key);
    let json_res : Value = serde_json::from_str(&res.unwrap()).unwrap();
    return json_res["id"].as_str().unwrap().to_string();
}

fn sp_get_memberids(key: &str, system_id: &str) -> HashMap<String, String> {
    let url = format!("{}/members/{}", SP_URL, system_id);

    let res = http_get(url,key);
    let datas: Vec<Value> = serde_json::from_str(&res.unwrap()).unwrap();

    let mut sp_memberdata: HashMap<String, String> = HashMap::new();
    for data in datas {
        sp_memberdata.insert(String::from(data["content"]["name"].as_str().unwrap()), String::from(data["id"].as_str().unwrap()));
    }

    return sp_memberdata;
}

fn get_sp_id(mem: &Member, ids: &HashMap<String, String>) -> String {

    let mut member_id = String::new();

    for (mn, mid) in ids {
        if &mem.name == mn || &mem.alias == mn {
            member_id = String::from(mid);
        }
    }

    return member_id;
}

fn sp_get_fronters(key: &str, sys: &System) -> Vec<Member> {
    let url = format!("{}/fronters", SP_URL);

    let res = http_get(url,key);
    let datas: Vec<Value> = serde_json::from_str(&res.unwrap()).unwrap();

    let mut members = Vec::new();
    for data in datas {
        let sp_id = &data["content"]["member"].as_str().unwrap();
        for member in &sys.members {
            if &member.sp_id == sp_id {
                members.push(member.clone());
            }
        }

    }
    return members;
}

fn sp_get_frontids(key: &str, to_id: &Member) -> String {
    let url = format!("{}/fronters", SP_URL);

    let res = http_get(url,key);
    let datas: Vec<Value> = serde_json::from_str(&res.unwrap()).unwrap();

    for data in datas {
        let sp_f_id = &data["id"].as_str().unwrap();
        let sp_id = &data["content"]["member"].as_str().unwrap();
        if sp_id.to_string() == to_id.sp_id {
            return sp_f_id.to_string();
        }
    }
    return String::new();
}

fn sp_set_fronters(key: &str, to_front: &Vec<Member>, fronters: &Fronters) {
    if to_front == &fronters.sp {
        println!("Members already fonting");
        return;
    }

    for fronting_member in &fronters.sp {
        if !to_front.contains(&fronting_member) {
            let f_id = sp_get_frontids(&key, &fronting_member);

            let url = format!("{}/frontHistory/{}", SP_URL, f_id);

            #[derive (Serialize)]
            #[allow (non_snake_case)]
            struct SpRem {
                live: bool,
                endTime: u64
            }

            let end_time = std::time::SystemTime::now().duration_since(std::time::SystemTime::UNIX_EPOCH).expect("wa").as_secs();
            let rem = SpRem {
                live: false,
                endTime: end_time * 1000
            };

            let body = serde_json::to_string(&rem).expect("Error");


            let client = reqwest::Client::new();
            let rb = client
                .patch(url)
                .body(body)
                .header("content-type", "application/json; charset=utf-8")
                .header(USER_AGENT, "Pluralsync")
                .header(AUTHORIZATION, key);

            match http_request(rb) {
                Ok(_) => (),
                Err(e) => println!("{}", e.to_string()),
            }
        }
    }

    for tf_member in to_front {
        if !fronters.sp.contains(&tf_member) {
            let url = format!("{}/frontHistory", SP_URL);

            #[derive (Serialize)]
            #[allow (non_snake_case)]
            struct SpAdd {
                member: String,
                custom: bool,
                live: bool,
                startTime: u64
            }

            let start_time = std::time::SystemTime::now().duration_since(std::time::SystemTime::UNIX_EPOCH).expect("wa").as_secs();
            let rem = SpAdd {
                member: String::from(&tf_member.sp_id),
                custom: false,
                live: true,
                startTime: start_time * 1000
            };

            let body = serde_json::to_string(&rem).expect("Error");

            let client = reqwest::Client::new();
            let rb = client
                .post(url)
                .body(body)
                .header("content-type", "application/json; charset=utf-8")
                .header(USER_AGENT, "Pluralsync")
                .header(AUTHORIZATION, key);

            match http_request(rb) {
                Ok(_) => (),
                Err(e) => println!("{}", e.to_string()),
            }
        }
    }
}

fn load_json(path: String) -> Value {
    if Path::new(&path).exists() {
        let config_data = fs::read_to_string(&path).expect("File not found");
        return serde_json::from_str(&config_data).unwrap();
    } else {
        println!("Config file in {path} not found");
        return Value::Null;
    }
}

fn get_config(config_path: &str) -> Result<Config, &'static str> {
    let path = format!("{}/config.json", config_path);
    if Path::new(config_path).exists() {

        let result = load_json(String::from(&path));

        if result == Value::Null {
            let _ = fs::write(path, EXAMPLE_JSON);
            return Err("Config file missing, creating template in {path}");
        } else {
            let config: Config = serde_json::from_value(result).expect("Error unwrapping");
            return Ok(config);
        }
    } else {
        let _ = create_dir(config_path);
        let _ = fs::write(path, EXAMPLE_JSON);
        return Err("Directory {config_path} does not exist. Creating with template config");
    }
}

fn get_system(config_path: &str) -> System {
    let path = format!("{}/system.json", config_path);

    let mut result = load_json(String::from(&path));

    if result == Value::Null {
        println!("Syncing system config");
        let _ = sync(String::from(config_path));
        result = load_json(String::from(&path));
    }

    let vec = serde_json::to_vec(&result).unwrap();
    let sys = serde_json::from_slice::<System>(&vec).unwrap();

    return sys;

}

fn get_fronters(pk_key: &str, sp_key: &str, sys: &System, ff: ForceFrom) -> Fronters {
    let mut fronters = Fronters::new(
        pk_get_fronters(pk_key, sys),
        sp_get_fronters(sp_key, sys)
    );

    if fronters.pk != fronters.sp {
        match ff {
            ForceFrom::PK => {
                fronters.sp = fronters.pk.clone()
            }
            ForceFrom::SP => {
                fronters.pk = fronters.sp.clone()
            }
            ForceFrom::None => {
            }
        }
    }

    return fronters;
}

#[tokio::main]
async fn http_get(url: String, key: &str) -> Result<String, Box<dyn std::error::Error>> {
    let client = reqwest::Client::new();
    let rb = client
        .get(url)
        .header(USER_AGENT, "Pluralsync")
        .header(AUTHORIZATION, key);
    let res = rb.send()
        .await?
        .text()
        .await?;

    Ok(res)
}

#[tokio::main]
async fn http_request(rb: RequestBuilder) -> Result<String, Box<dyn std::error::Error>> {
    let res = rb.send()
        .await?
        .text()
        .await?;
    Ok(res)
}