3 files changed, 71 insertions, 89 deletions

diff --git a/src/main.rs b/src/main.rs
index b12f5a6..7570dcb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,63 +3,13 @@ use std::io::prelude::*;
 
 extern crate notify_rust;
 extern crate toml;
-extern crate websocket;
-extern crate rustc_serialize;
 use notify_rust::Notification;
 
-mod message;
+pub mod message;
 use message::{Message, Push};
 
-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 process_message(message: Result<websocket::Message, websocket::result::WebSocketError>) -> Option<Message> {
-        let message = match message {
-            Ok(m) => m,
-            Err(e) => {
-                println!("Error: {:?}", e);
-                return None;
-            }
-        };
-
-        if let websocket::Message::Text(message) = message {
-                Message::parse(message.as_ref())
-        }
-        else {
-            None
-        }
-    }
-
-    fn messages<'a>(&'a mut self) -> Box<Iterator<Item=Message> + 'a> {
-        let mut receiver = self.client.get_mut_reciever(); // there is a typo in the API
-        Box::new(receiver.incoming_messages().filter_map(PBClient::process_message))
-    }
-}
+pub mod pushbullet;
+use pushbullet::PBClient;
 
 fn main() {
     let mut cfg_file = File::open("config.toml").expect("Could not find config.toml.");
@@ -76,45 +26,23 @@ fn main() {
         .expect("'token' should be a string");
 
     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 {
-            Ok(m) => m,
-            Err(e) => {
-                println!("Error: {:?}", e);
-                return;
-            }
-        };
 
+    for message in client.messages() {
         match message {
-            websocket::Message::Close(_) => {
-                return;
-            }
-            websocket::Message::Text(message) => {
-                let msg = Message::parse(message.as_ref());
-                if let Some(msg) = msg {
-                    match msg {
-                        Message::Push(Push::Mirror {
-                            title,
-                            body,
-                            application_name,
-                            ..
-                        }) => {
-                            let title = format!("{}: {}", application_name, title);
-                            Notification::new()
-                                .body(body.as_ref())
-                                .summary(title.as_ref())
-                                .show()
-                                .unwrap();
-                        }
-                        _ => {}
-                    }
-                }
-            }
-            _ => {
-                println!("Got {:?}", message);
+            Message::Push(Push::Mirror {
+                title,
+                body,
+                application_name,
+                ..
+            }) => {
+                let title = format!("{}: {}", application_name, title);
+                Notification::new()
+                    .body(body.as_ref())
+                    .summary(title.as_ref())
+                    .show()
+                    .unwrap();
             }
+            _ => {}
         }
     }
 }
diff --git a/src/message.rs b/src/message/mod.rs
index 0617c4d..0617c4d 100644
--- a/src/message.rs
+++ b/src/message/mod.rs
diff --git a/src/pushbullet/mod.rs b/src/pushbullet/mod.rs
new file mode 100644
index 0000000..1f2254e
--- /dev/null
+++ b/src/pushbullet/mod.rs
@@ -0,0 +1,54 @@
+extern crate websocket;
+extern crate rustc_serialize;
+
+use message::Message;
+
+use self::websocket::Sender;
+use self::websocket::Receiver;
+use self::websocket::WebSocketStream;
+use self::websocket::DataFrame;
+use self::websocket::Client;
+use self::websocket::client::request::Url;
+
+pub 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 {
+    pub fn new(token: &str) -> PBClient {
+
+        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 process_message(message: Result<websocket::Message, websocket::result::WebSocketError>) -> Option<Message> {
+        let message = match message {
+            Ok(m) => m,
+            Err(e) => {
+                println!("Error: {:?}", e);
+                return None;
+            }
+        };
+
+        if let websocket::Message::Text(message) = message {
+                Message::parse(message.as_ref())
+        }
+        else {
+            None
+        }
+    }
+
+    pub fn messages<'a>(&'a mut self) -> Box<Iterator<Item=Message> + 'a> {
+        let mut receiver = self.client.get_mut_reciever(); // there is a typo in the API
+        Box::new(receiver.incoming_messages().filter_map(PBClient::process_message))
+    }
+}