1 files changed, 105 insertions, 25 deletions

diff --git a/src/main.rs b/src/main.rs
index 315d093..e96bbd7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,26 +1,31 @@
 
 mod protos;
 mod fancy;
-use protos::protos::gtfs_realtime::FeedMessage;
+use protos::protos::gtfs_realtime::{FeedMessage, TripUpdate, VehiclePosition, TripDescriptor, EntitySelector, Alert, FeedEntity};
+use std::io::Read;
 
 use protobuf::Message;
 use clap::Parser;
 // use anyhow::Context;
 use miette::{WrapErr, IntoDiagnostic, miette};
-use reqwest::header::{HeaderValue, CONTENT_TYPE};
+use reqwest::{header::{HeaderValue, CONTENT_TYPE}, Url};
 
 #[derive(Parser, Debug)]
 #[command(author, version, about, long_about = None)]
 struct Args {
-   /// uri of the GTFS RT feed to fetch & display
-   url: String,
+   /// Either a file or an URI containing the gtfs realtime feed.
+   /// Omit this to read from stdin instead.
+   url_or_file: Option<String>,
    /// emit the feed as json
    #[arg(long)]
    json: bool,
-   /// ignore things that look wrong as long as possible
-   #[arg(long="ignore-nonfatal", short='i')]
-   ignore_nonfatal: bool,
-   /// don't do terminal colours
+   /// filter for FeedEntities affecting the given trip id
+   #[arg(short='f', long="filter-trip")]
+   filter_trip: Option<String>,
+   /// ignore the Content-Type header (default is to abort if it's not application/octet-string)
+   #[arg(long="ignore-content-type", short='i')]
+   ignore_content_type: bool,
+   /// don't do terminal colours in output
    #[arg(long="no-colors")]
    no_colors: bool
 }
@@ -30,29 +35,60 @@ struct Args {
 async fn main() -> miette::Result<()> {
 
     let args = Args::parse();
-    let resp = reqwest::get(&args.url)
-        .await.into_diagnostic().wrap_err("Request failed")?
-        .error_for_status().into_diagnostic()?;
-
-    if !args.ignore_nonfatal {
-        let ct = HeaderValue::from_static("application/octet-stream");
-        match resp.headers().get(CONTENT_TYPE) {
-            Some(content_type) if ct == content_type => (),
-            Some(other_type) =>
-                Err(miette!("should be {:?}", ct))
-                  .wrap_err(format!("Bad Content-Type {:?}", other_type))?,
-            None =>
-                Err(miette!("Content-Type header is missing"))?
+
+    let resp = match args.url_or_file {
+        Some(url_or_file) => match Url::parse(&url_or_file) {
+            Ok(url) => {
+                let resp = reqwest::get(url).await
+                    .into_diagnostic()
+                    .wrap_err("Request failed")?
+                    .error_for_status().into_diagnostic()?;
+
+                if !args.ignore_content_type {
+                    let ct = HeaderValue::from_static("application/octet-stream");
+                    match resp.headers().get(CONTENT_TYPE) {
+                        Some(content_type) if ct == content_type => (),
+                        Some(other_type) =>
+                            Err(miette!("should be {:?}", ct))
+                            .wrap_err(format!("Bad Content-Type {:?}", other_type))?,
+                        None =>
+                            Err(miette!("Content-Type header is missing"))?
+                    }
+                }
+
+                resp.bytes().await.into_diagnostic()?.into()
+            },
+            Err(_) => {
+                std::fs::read(&url_or_file)
+                    .into_diagnostic()
+                    .wrap_err(format!("failed to read file: {}", url_or_file))?
+            }
+        },
+        None => {
+            let mut buf = Vec::new();
+            let mut stdin = std::io::stdin();
+            stdin.read_to_end(&mut buf).into_diagnostic()?;
+            buf
         }
-    }
+    };
 
-    let resp = resp
-        .bytes().await.into_diagnostic()?;
 
-    let proto = FeedMessage::parse_from_bytes(&resp[..])
+    let mut proto = FeedMessage::parse_from_bytes(&resp[..])
         .into_diagnostic()
         .wrap_err("Could not parse protobuf format")?;
 
+    if let Some(filter) = args.filter_trip {
+        proto = FeedMessage {
+            entity: proto
+                .entity
+                .iter()
+                .filter(|entity| entity.concerns(&filter))
+                .map(|s| s.to_owned())
+                .collect(),
+            ..proto
+        };
+    }
+
     match args.json {
         true => 
             println!("{}", protobuf_json_mapping::print_to_string(&proto).into_diagnostic()?),
@@ -65,3 +101,47 @@ async fn main() -> miette::Result<()> {
     Ok(())
 }
 
+trait Concerns {
+    fn concerns(&self, trip_id: &str) -> bool;
+}
+
+impl Concerns for TripDescriptor {
+    fn concerns(&self, trip_id: &str) -> bool {
+        match &self.trip_id {
+            Some(id) if id == trip_id => true,
+            _ => false
+        }
+    }
+}
+
+impl Concerns for TripUpdate {
+    fn concerns(&self, trip_id: &str) -> bool {
+        self.trip.concerns(trip_id)
+    }
+}
+
+impl Concerns for VehiclePosition {
+    fn concerns(&self, trip_id: &str) -> bool {
+        self.trip.concerns(trip_id)
+    }
+}
+
+impl Concerns for EntitySelector {
+    fn concerns(&self, trip_id: &str) -> bool {
+        self.trip.concerns(trip_id)
+    }
+}
+
+impl Concerns for Alert {
+    fn concerns(&self, trip_id: &str) -> bool {
+        self.informed_entity.iter().any(|e| e.concerns(trip_id))
+    }
+}
+
+impl Concerns for FeedEntity {
+    fn concerns(&self, trip_id: &str) -> bool {
+        self.trip_update.concerns(trip_id)
+            || self.vehicle.concerns(trip_id)
+            || self.alert.concerns(trip_id)
+    }
+}