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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
|
/*
* einhyrningsinsctl: controller/shell for einhyrningsins
* Copyright (C) 2016 Bryan Newbold <bnewbold@robocracy.org>
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
#[macro_use]
extern crate json;
extern crate getopts;
extern crate log;
extern crate env_logger;
extern crate nix;
extern crate timer;
extern crate time;
extern crate chan_signal;
extern crate url;
extern crate rustyline;
use std::io;
use std::io::prelude::*;
use std::io::{BufReader, BufWriter};
use std::env;
use std::path::Path;
use std::process::exit;
use std::os::unix::net::UnixStream;
use getopts::Options;
use rustyline::error::ReadlineError;
use rustyline::Editor;
// This is the main event loop
fn shell(ctrl_stream: UnixStream) {
let mut reader = BufReader::new(&ctrl_stream);
let mut writer = BufWriter::new(&ctrl_stream);
// `()` can be used when no completer is required
let mut rl = Editor::<()>::new();
println!("");
println!("Welcome to the einhyrningsins shell!");
println!("Try 'help if you need it");
loop {
let readline = rl.readline("> ");
match readline {
Ok(line) => {
rl.add_history_entry(&line);
if line.len() == 0 { continue; };
let mut chunks = line.split(' ');
let cmd = chunks.nth(0).unwrap();
let args = chunks.collect();
match send_msg(&mut reader, &mut writer, cmd, args) {
Ok(s) => { println!("{}", s); },
Err(e) => {
println!("Error sending control message: {}", e);
exit(-1);
},
}
},
Err(ReadlineError::Interrupted) | Err(ReadlineError::Eof) => {
println!("Caught kill signal (shutting down)");
break
},
Err(err) => {
println!("Shell Error: {:?} (shutting down)", err);
break
}
}
}
//drop(ctrl_stream);
}
// This function sends a single request message down the writer, then waits for a reply on the
// reader and prints the result.
fn send_msg(reader: &mut BufRead, writer: &mut Write, cmd: &str, args: Vec<&str>) -> io::Result<String> {
let mut buffer = String::new();
let mut arg_list = json::JsonValue::new_array();
for a in args {
arg_list.push(a).expect("function args");
}
let req = object!{
"command" => cmd,
"args" => arg_list
};
//println!("Sending: {}", req.dump());
try!(writer.write_all(format!("{}\n", req.dump()).as_bytes()));
try!(writer.flush());
try!(reader.read_line(&mut buffer));
//println!("Got: {}", buffer);
let reply = match json::parse(&buffer) {
Ok(obj) => obj,
Err(_) => { return Ok(buffer) },
};
Ok(match reply.as_str() {
Some(s) => s.to_string(),
None => buffer,
})
}
fn print_usage(opts: Options) {
let brief = "usage:\teinhyrningsinsctl [options] program";
println!("");
print!("{}", opts.usage(&brief));
}
fn main() {
let args: Vec<String> = env::args().collect();
let mut opts = Options::new();
opts.optflag("h", "help", "print this help menu");
opts.optflag("", "version", "print the version");
opts.optopt("e", "execute", "submit this command instead (no shell)", "CMD");
opts.optopt("d", "socket-path", "where to look for control socket (default: /tmp/einhorn.sock)", "PATH");
let matches = match opts.parse(&args[1..]) {
Ok(m) => { m }
Err(f) => { println!("{}", f.to_string()); print_usage(opts); exit(-1); }
};
if matches.opt_present("help") {
print_usage(opts);
return;
}
if matches.opt_present("version") {
println!("einhyrningsinsctl {}", env!("CARGO_PKG_VERSION"));
return;
}
// Bind to Control Socket
let path_str = matches.opt_str("socket-path").unwrap_or("/tmp/einhorn.sock".to_string());
let ctrl_path = Path::new(&path_str);
if !ctrl_path.exists() {
println!("Couldn't find control socket ({:?})", ctrl_path);
println!("Is the master process running? Do you need to tell me the correct socket path?");
exit(-1);
}
//println!("Connecting to control socket: {:?}", ctrl_path);
let ctrl_stream = match UnixStream::connect(ctrl_path) {
Ok(s) => s,
Err(e) => {
println!("Couldn't open socket [{}]: {}", path_str, e);
exit(-1);
},
};
// Send a test message before continuing
send_msg(&mut BufReader::new(&ctrl_stream),
&mut BufWriter::new(&ctrl_stream),
"ehlo",
vec![]).unwrap();
match matches.opt_str("execute") {
Some(cmd) => {
match send_msg(&mut BufReader::new(&ctrl_stream),
&mut BufWriter::new(&ctrl_stream),
&cmd, vec![]) {
Ok(reply) => println!("{}", reply),
Err(e) => println!("Communications error: {}", e),
}
},
None => shell(ctrl_stream),
}
exit(0);
}
|
