plankircd/src/models.rs

use std::fmt;

#[derive(Debug, PartialEq, Eq)]
pub enum IrcCommand {
  USER,
  NICK,
  JOIN,
  PART,
  PRIVMSG,
  WHO,
  MODE,
  LIST,
  CAP,
  PASS,
  QUIT,
  PING,
  PONG,
  ERROR,
  Num(u16),
}

impl IrcCommand {
  fn from_str(s: &str) -> Option<IrcCommand> {
    use self::IrcCommand::*;
    match s {
      "USER" => Some(USER),
      "NICK" => Some(NICK),
      "JOIN" => Some(JOIN),
      "PART" => Some(PART),
      "PRIVMSG" => Some(PRIVMSG),
      "WHO" => Some(WHO),
      "MODE" => Some(MODE),
      "CAP" => Some(CAP),
      "QUIT" => Some(QUIT),
      "LIST" => Some(LIST),
      "PASS" => Some(PASS),
      "PING" => Some(PING),
      "PONG" => Some(PONG),
      "ERROR" => Some(ERROR),
      s => if let Ok(i) = s.parse::<u16>() {
        Some(Num(i))
      } else { None },
    }
  }
}

impl fmt::Display for IrcCommand {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    use self::IrcCommand::*;
    match self {
      Num(i) => write!(f, "{:03}", i),
      cmd => write!(f, "{:?}", cmd),
    }
  }
}

#[derive(Debug, PartialEq, Eq)]
pub struct IrcMessage {
  pub prefix: Option<String>,
  pub command: IrcCommand,
  pub params: Vec<String>,
}

impl IrcMessage {
  pub fn new(pre: Option<&str>, cmd: IrcCommand, prm: Vec<&str>) -> IrcMessage {
    IrcMessage {
      prefix: if let Some(s) = pre { Some(s.to_string()) } else { None },
      command: cmd,
      params: prm.iter().map(|s| s.to_string()).collect(),
    }
  }

  pub fn from_str(msg: &str) -> Result<IrcMessage, String> { // TODO: Wrap in Result<>
    let mut args = msg.trim().splitn(2, " :");
    let mut left = args.next().unwrap().split(' ');

    let mut prefix = None;
    let mut cmd = None;

    if let Some(s) = left.next() {
      if s.starts_with(':') {
        prefix = Some(&s[1..]);

        if let Some(c) = left.next() {
          cmd = match IrcCommand::from_str(c) {
            Some(c) => Some(c),
            None => return Err(format!("Unknown command: {}", s)),
          };
        };
      }
      else {
        cmd = match IrcCommand::from_str(s) {
          Some(c) => Some(c),
          None => return Err(format!("Unknown command: {}", s)),
        };
      }
    }

    let mut params: Vec<&str> = left.collect();

    if let Some(s) = args.next() {
      params.push(s);
    };

    match cmd {
      Some(c) => Ok(IrcMessage::new(prefix, c, params)),
      None => Err("Not enough arguments".to_string()),
    }
  }

  pub fn to_string(&self) -> String {
    let param: Vec<String> = self.params.iter()
      .map(|s| if s.contains(" ") { format!(":{}", s) } else { s.clone() }) 
      .collect();
    // TODO: get rid of clone, also would be better to only check and (...)
    // replace the last element since we assume that the preceding params
    // doesn't have any spaces

    match self.prefix {
      Some(ref prefix) => format!(":{} {} {}\r\n", prefix, self.command, param.join(" ")),
      None => format!("{} {}\r\n", self.command, param.join(" "))
    }
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn message_from_string() {
    let msg = IrcMessage::from_str(":tlx PRIVMSG supes :u suk lol\r\n").unwrap();

    assert_eq!(msg, IrcMessage {
      prefix: Some("tlx".to_string()),
      command: IrcCommand::PRIVMSG,
      params: vec!["supes".to_string(), "u suk lol".to_string()]
    });
  }

  #[test]
  fn message_from_string_with_multiple_colons_in_trailing() {
    let msg = IrcMessage::from_str(":tlx PRIVMSG supes :u suk lol :D\r\n").unwrap();

    assert_eq!(msg, IrcMessage {
      prefix: Some("tlx".to_string()),
      command: IrcCommand::PRIVMSG,
      params: vec!["supes".to_string(), "u suk lol :D".to_string()]
    });
  }

  #[test]
  fn message_from_string_without_prefix() {
    let msg = IrcMessage::from_str("PRIVMSG supes :u suk lol\r\n").unwrap();

    assert_eq!(msg, IrcMessage {
      prefix: None,
      command: IrcCommand::PRIVMSG,
      params: vec!["supes".to_string(), "u suk lol".to_string()]
    });
  }

  #[test]
  fn message_to_string() {
    let msg = IrcMessage {
      prefix: Some("tlx".to_string()),
      command: IrcCommand::PRIVMSG,
      params: vec!["supes".to_string(), "u suk lol".to_string()]
    };

    assert_eq!(&msg.to_string(), ":tlx PRIVMSG supes :u suk lol\r\n");
  }

  #[test]
  fn message_to_string_without_prefix() {
    let msg = IrcMessage {
      prefix: None,
      command: IrcCommand::PRIVMSG,
      params: vec!["supes".to_string(), "u suk lol".to_string()]
    };

    assert_eq!(&msg.to_string(), "PRIVMSG supes :u suk lol\r\n");
  }

  #[test]
  fn message_to_string_numeric() {
    let msg = IrcMessage {
      prefix: Some("plankircd".to_string()),
      command: IrcCommand::Num(4),
      params: vec!["supes".to_string(), "u suk lol".to_string()]
    };

    assert_eq!(&msg.to_string(), ":plankircd 004 supes :u suk lol\r\n");
  }
}