1 files changed, 30 insertions, 11 deletions

diff --git a/src/main.rs b/src/main.rs
index d904f6a..ea58059 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,11 +10,35 @@ use notify_rust::Notification;
 mod message;
 use message::{Message, Push};
 
-fn main() {
-    use websocket::{Receiver};
-    use websocket::client::request::Url;
-    use websocket::Client;
+use websocket::Sender;
+use websocket::Receiver;
+use websocket::WebSocketStream;
+use websocket::DataFrame;
+use websocket::Client;
+struct PBClient {
+    client: websocket::client::Client<websocket::dataframe::DataFrame,
+                                      websocket::client::sender::Sender<websocket::stream::WebSocketStream>,
+                                      websocket::client::receiver::Receiver<websocket::stream::WebSocketStream>>
+}
+
+impl PBClient {
+    fn new(token: &str) -> PBClient {
+        use websocket::client::request::Url;
+        use websocket::Client;
+
+        let wss_url = format!("wss://stream.pushbullet.com/websocket/{}", token);
+        let url = Url::parse(wss_url.as_ref()).unwrap();
+        let request = Client::connect(url).unwrap();
+        let response = request.send().unwrap();
+        response.validate().unwrap();
+
+        PBClient {
+            client: response.begin()
+        }
+    }
+}
 
+fn main() {
     let mut cfg_file = File::open("config.toml").expect("Could not find config.toml.");
     let mut s = String::new();
     cfg_file.read_to_string(&mut s).expect("Could not read config.toml");
@@ -28,13 +52,8 @@ fn main() {
         .as_str()
         .expect("'token' should be a string");
 
-    let wss_url = format!("wss://stream.pushbullet.com/websocket/{}", token);
-    let url = Url::parse(wss_url.as_ref()).unwrap();
-    let request = Client::connect(url).unwrap();
-    let response = request.send().unwrap();
-    response.validate().unwrap();
-
-    let (mut sender, mut receiver) = response.begin().split();
+    let mut client = PBClient::new(token.as_ref());
+    let mut receiver = client.client.get_mut_reciever();
 
     for message in receiver.incoming_messages() {
         let message = match message {