bie 0.1.0

Drop a file from a remote without almost any access
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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

mod settings;
mod protocol;

use anyhow::Result;
use clap::{Parser, Subcommand};
use protocol::BieProtocol;
use std::io::Write;
use std::path::PathBuf;
use tempfile::NamedTempFile;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::protocol::Message;
use url::Url;

use futures::{SinkExt, StreamExt};

#[derive(Parser)]
#[command(name = "bie")]
#[command(bin_name = "bie")]
struct Cli {
    /// Subcommand to execute
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Download a file
    Get {
        /// The name of the file to download
        file_name: PathBuf,
    },

    /// Show configuration
    Config,
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    // Parse command line arguments
    let cli = Cli::parse();

    match cli.command {
        Commands::Config => {
            let settings = settings::Settings::load()?;
            println!("{:?}", settings);
            Ok(())
        }
        Commands::Get { file_name } => {
            let settings = settings::Settings::load()?;

            // First of all - we need to connect to the server via websocket and request a token
            let url = create_websocket_url(&settings.bastion_server_url)?;
            let ws_url = url.to_string();
            let (ws_stream, _) = connect_async(&ws_url).await?;

            let (mut write, mut read) = ws_stream.split();

            // The first message from the server is a token
            let token_message = read.next().await;
            let token_message =
                token_message.ok_or_else(|| anyhow::anyhow!("No token message received"))?;
            let token_message = token_message?;

            let token = match token_message {
                Message::Text(token) => token,
                _ => {
                    return Err(anyhow::anyhow!("Invalid token message"));
                }
            };

            // Generate the link for curl
            let url = format!("{}/upload/{}", settings.bastion_server_url, token);
            println!("In order to send a file, use this snippet:\n");
            println!("echo <your-file-name> | xargs -I{{}} curl -v -X POST -F \"file=@{{}}\" {}\n", url);
            println!("<type echo then <tab> to get file name autocompletion, then copy snippet from \"|\", or Ctrl-C to exit>");

            // Here we start a loop to write all incoming content into a file

            // First - create temp file in tmp directory
            let mut temp_file = NamedTempFile::new()?;

            // Now - start receiving loop for websocket messages
            loop {
                let message = read.next().await;
                let message = message.ok_or_else(|| anyhow::anyhow!("No message received"))?;
                let message = message?;

                match message {
                    Message::Binary(data) => {
                        // Here we should parse CBOR encoded message
                        match BieProtocol::from(&data[..]) {
                            BieProtocol::FileChunk(chunk) => {
                                temp_file.write_all(&chunk)?;
                            }
                            BieProtocol::EndOfFile => {
                                break;
                            }
                            _ => {
                                return Err(anyhow::anyhow!("Invalid message"));
                            }
                        }
                    }
                    Message::Close(_) => {
                        return Ok(());
                    }
                    _ => {
                        return Err(anyhow::anyhow!("Invalid message"));
                    }
                }
            }
            // Here we need to flush the tempfile and copy it to right place
            temp_file.flush()?;
            temp_file.persist(file_name.as_path())?;
            println!("File saved to {}", file_name.display());

            write.close().await?;

            Ok(())
        }
    }
}

fn create_websocket_url(server_url: &str) -> Result<Url, anyhow::Error> {
    let mut url = Url::parse(server_url)?;
    // Here the scheme is valid
    if url.scheme() == "http" {
        url.set_scheme("ws")
            .map_err(|_| anyhow::anyhow!("Invalid scheme"))?;
    } else if url.scheme() == "https" {
        url.set_scheme("wss")
            .map_err(|_| anyhow::anyhow!("Invalid scheme"))?;
    } else {
        return Err(anyhow::anyhow!("Invalid scheme"));
    }

    url.set_path("/wait_file");
    Ok(url)
}