iptoasn-webservice/src/webservice.rs

use asns::*;
use horrorshow::prelude::*;
use iron::{BeforeMiddleware, typemap};
use iron::headers::{CacheControl, CacheDirective, Accept};
use iron::mime::*;
use iron::modifiers::Header;
use iron::prelude::*;
use iron::status;
use router::Router;
use serde_json;
use std::net::IpAddr;
use std::str::FromStr;
use std::sync::{Arc, RwLock};

const TTL: u32 = 86400;

struct ASNsMiddleware {
    asns_arc: Arc<RwLock<Arc<ASNs>>>,
}

impl typemap::Key for ASNsMiddleware {
    type Value = Arc<ASNs>;
}

impl ASNsMiddleware {
    fn new(asns_arc: Arc<RwLock<Arc<ASNs>>>) -> ASNsMiddleware {
        ASNsMiddleware { asns_arc: asns_arc }
    }
}

impl BeforeMiddleware for ASNsMiddleware {
    fn before(&self, req: &mut Request) -> IronResult<()> {
        req.extensions
            .insert::<ASNsMiddleware>(self.asns_arc.read().unwrap().clone());
        Ok(())
    }
}

enum OutputType {
    Json,
    Html,
}

pub struct WebService;

impl WebService {
    fn index(_: &mut Request) -> IronResult<Response> {
        Ok(Response::with((status::Ok,
                           Mime(TopLevel::Text,
                                SubLevel::Plain,
                                vec![(Attr::Charset, Value::Utf8)]),
                           Header(CacheControl(vec![CacheDirective::Public,
                                                    CacheDirective::MaxAge(TTL)])),
                           "See https://iptoasn.com")))
    }

    fn accept_type(req: &Request) -> OutputType {
        let mut output_type = OutputType::Json;
        if let Some(header_accept) = req.headers.get::<Accept>() {
            for header in header_accept.iter() {
                match header.item {
                    Mime(TopLevel::Text, SubLevel::Html, _) => {
                        output_type = OutputType::Html;
                        break;
                    }
                    Mime(_, SubLevel::Json, _) => {
                        output_type = OutputType::Json;
                        break;
                    }
                    _ => {}
                }
            }
        }
        output_type
    }

    fn output_json(_ip_str: &str,
                   map: serde_json::Map<String, serde_json::value::Value>,
                   cache_header: Header<CacheControl>)
                   -> IronResult<Response> {
        let json = serde_json::to_string(&map).unwrap();
        let mime_json = Mime(TopLevel::Application,
                             SubLevel::Json,
                             vec![(Attr::Charset, Value::Utf8)]);
        Ok(Response::with((status::Ok, mime_json, cache_header, json)))
    }

    fn output_html(ip_str: &str,
                   map: serde_json::Map<String, serde_json::value::Value>,
                   cache_header: Header<CacheControl>)
                   -> IronResult<Response> {
        let mime_html = Mime(TopLevel::Text,
                             SubLevel::Html,
                             vec![(Attr::Charset, Value::Utf8)]);
        let html = html!{
            head {
                title { : "iptoasn lookup" }
                meta(name="viewport", content="width=device-widthinitial-scale=1");
                link(rel="stylesheet", href="https://maxcdn.bootstrapcdn.com/bootstrap/4.0.0-alpha.5/css/bootstrap.min.css", integrity="sha384-AysaV+vQoT3kOAXZkl02PThvDr8HYKPZhNT5h/CXfBThSRXQ6jW5DO2ekP5ViFdi", crossorigin="anonymous");
                style {
                    : "body { margin: 1em 4em }"
                }
            }
            body(class="container-fluid") {
                header {
                    h1 { : format_args!("Information for IP address: {}", ip_str) }
                }
                table {
                    tr {
                        th { : "Announced" }
                        td { : format_args!("{}", if map.get("announced")
                            .unwrap().as_bool().unwrap() { "Yes" } else { "No" }) }
                    }
                    @ if map.get("announced").unwrap().as_bool().unwrap() == true {
                        tr {
                            th { : "First IP" }
                            td { : format_args!("{}", map.get("first_ip")
                                .unwrap().as_str().unwrap()) }
                        }
                        tr {
                            th { : "Last IP" }
                            td { : format_args!("{}", map.get("last_ip")
                                .unwrap().as_str().unwrap()) }
                        }
                        tr {
                            th { : "AS Number" }
                            td { : format_args!("{}", map.get("as_number")
                                .unwrap().as_u64().unwrap()) }
                        }
                        tr {
                            th { : "AS Country code" }
                            td { : format_args!("{}", map.get("as_country_code")
                                .unwrap().as_str().unwrap()) }
                        }
                        tr {
                            th { : "AS Description" }
                            td { : format_args!("{}", map.get("as_description")
                                .unwrap().as_str().unwrap()) }
                        }
                    }
                }
            }
        }
            .into_string()
            .unwrap();
        let html = format!("<!DOCTYPE html>\n<html>{}</html>", html);
        Ok(Response::with((status::Ok, mime_html, cache_header, html)))
    }

    fn output(output_type: OutputType,
              ip_str: &str,
              map: serde_json::Map<String, serde_json::value::Value>,
              cache_header: Header<CacheControl>)
              -> IronResult<Response> {
        match output_type {
            OutputType::Json => Self::output_json(ip_str, map, cache_header),
            _ => Self::output_html(ip_str, map, cache_header),
        }
    }

    fn ip_lookup(req: &mut Request) -> IronResult<Response> {
        let mime_text = Mime(TopLevel::Text,
                             SubLevel::Plain,
                             vec![(Attr::Charset, Value::Utf8)]);
        let cache_header = Header(CacheControl(vec![CacheDirective::Public,
                                                    CacheDirective::MaxAge(TTL)]));
        let ip_str = match req.extensions.get::<Router>().unwrap().find("ip") {
            None => {
                let response = Response::with((status::BadRequest,
                                               mime_text,
                                               cache_header,
                                               "Missing IP address"));
                return Ok(response);
            }
            Some(ip_str) => ip_str,
        };
        let ip = match IpAddr::from_str(ip_str) {
            Err(_) => {
                return Ok(Response::with((status::BadRequest,
                                          mime_text,
                                          cache_header,
                                          "Invalid IP address")));
            }
            Ok(ip) => ip,
        };
        let asns = req.extensions.get::<ASNsMiddleware>().unwrap();
        let found = match asns.lookup_by_ip(ip) {
            None => {
                let mut map = serde_json::Map::new();
                map.insert("announced".to_string(),
                           serde_json::value::Value::Bool(false));
                return Self::output(Self::accept_type(&req), ip_str, map, cache_header);
            }
            Some(found) => found,
        };
        let mut map = serde_json::Map::new();
        map.insert("announced".to_string(),
                   serde_json::value::Value::Bool(true));
        map.insert("first_ip".to_string(),
                   serde_json::value::Value::String(found.first_ip.to_string()));
        map.insert("last_ip".to_string(),
                   serde_json::value::Value::String(found.last_ip.to_string()));
        map.insert("as_number".to_string(),
                   serde_json::value::Value::Number(serde_json::Number::from(found.number)));
        map.insert("as_country_code".to_string(),
                   serde_json::value::Value::String(found.country.clone()));
        map.insert("as_description".to_string(),
                   serde_json::value::Value::String(found.description.clone()));
        Self::output(Self::accept_type(&req), ip_str, map, cache_header)
    }

    pub fn start(asns_arc: Arc<RwLock<Arc<ASNs>>>, listen_addr: &str) {
        let router = router!(index: get "/" => Self::index,
                             ip_lookup: get "/v1/as/ip/:ip" => Self::ip_lookup);
        let mut chain = Chain::new(router);
        let asns_middleware = ASNsMiddleware::new(asns_arc);
        chain.link_before(asns_middleware);
        warn!("webservice ready");
        Iron::new(chain).http(listen_addr).unwrap();
    }
}