Add strftime/strptime

kpcyrd · May 23, 2019 · b170145 · b170145
1 parent dc3f0f7
commit b170145
Show file tree

Hide file tree

Showing 4 changed files with 110 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -172,6 +172,9 @@ For everything else please have a look at the [detailed list][1].
   - [sock_newline](https://sn0int.readthedocs.io/en/latest/reference.html#sock-newline)
   - [status](https://sn0int.readthedocs.io/en/latest/reference.html#status)
   - [stdin_readline](https://sn0int.readthedocs.io/en/latest/reference.html#stdin-readline)
+  - [strftime](https://sn0int.readthedocs.io/en/latest/reference.html#strftime)
+  - [strptime](https://sn0int.readthedocs.io/en/latest/reference.html#strptime)
+  - [time_unix](https://sn0int.readthedocs.io/en/latest/reference.html#time-unix)
   - [url_decode](https://sn0int.readthedocs.io/en/latest/reference.html#url-decode)
   - [url_encode](https://sn0int.readthedocs.io/en/latest/reference.html#url-encode)
   - [url_escape](https://sn0int.readthedocs.io/en/latest/reference.html#url-escape)

diff --git a/docs/reference.rst b/docs/reference.rst
@@ -46,6 +46,10 @@ for ``DATETIME`` database fields.
 
     now = datetime()
 
+.. note::
+    This format is sn0int specific, to get the current time for scripting use
+    time_unix_ instead.
+
 db_add
 ------
 
@@ -741,6 +745,36 @@ Read a line from stdin. The final newline is not removed.
 .. note::
    This only works with `sn0int run --stdin`.
 
+strftime
+--------
+
+Format a timestamp generated with time_unix_ into a date, see `strftime rules`_.
+
+.. code-block:: lua
+
+    t = strftime('%d/%m/%Y %H:%M', 1558584994)
+
+strptime
+--------
+
+Parse a date into a unix timestamp, see `strftime rules`_.
+
+.. code-block:: lua
+
+    t = strptime('%d/%m/%Y %H:%M', '23/05/2019 04:16')
+
+.. _strftime rules: https://docs.rs/chrono/0.4.6/chrono/format/strftime/index.html
+
+time_unix
+---------
+
+Get the current time as seconds since ``January 1, 1970 0:00:00 UTC``, also
+known as UNIX timestamp. This timestamp can be formated using strftime_.
+
+.. code-block:: lua
+
+    now = time_unix()
+
 url_decode
 ----------
 

diff --git a/src/engine/ctx.rs b/src/engine/ctx.rs
@@ -383,6 +383,9 @@ fn ctx<'a>(env: Environment, logger: Arc<Mutex<Box<Reporter>>>) -> (hlua::Lua<'a
     runtime::sock_newline(&mut lua, state.clone());
     runtime::status(&mut lua, state.clone());
     runtime::stdin_readline(&mut lua, state.clone());
+    runtime::strftime(&mut lua, state.clone());
+    runtime::strptime(&mut lua, state.clone());
+    runtime::time_unix(&mut lua, state.clone());
     runtime::url_decode(&mut lua, state.clone());
     runtime::url_encode(&mut lua, state.clone());
     runtime::url_escape(&mut lua, state.clone());

diff --git a/src/runtime/datetime.rs b/src/runtime/datetime.rs
@@ -1,9 +1,11 @@
+use crate::errors::*;
 use crate::engine::ctx::State;
 use crate::hlua;
-use chrono::{Utc};
+use chrono::{NaiveDateTime, Utc};
 use std::sync::Arc;
 
 
+// TODO: consider renaming to time_sn0int
 pub fn datetime(lua: &mut hlua::Lua, _: Arc<State>) {
     lua.set("datetime", hlua::function0(move || -> String {
         let now = Utc::now().naive_utc();
@@ -12,6 +14,30 @@ pub fn datetime(lua: &mut hlua::Lua, _: Arc<State>) {
     }))
 }
 
+pub fn strftime(lua: &mut hlua::Lua, _: Arc<State>) {
+    lua.set("strftime", hlua::function2(move |format: String, time: i32| -> Result<String> {
+        let time = NaiveDateTime::from_timestamp_opt(time.into(), 0)
+            .ok_or_else(|| format_err!(""))?;
+        let time = time.format(&format).to_string();
+        Ok(time)
+    }))
+}
+
+pub fn strptime(lua: &mut hlua::Lua, state: Arc<State>) {
+    lua.set("strptime", hlua::function2(move |format: String, time: String| -> Result<i32> {
+        let datetime = NaiveDateTime::parse_from_str(&time, &format)
+            .map_err(|err| state.set_error(Error::from(err)))?;
+        Ok(datetime.timestamp() as i32)
+    }))
+}
+
+pub fn time_unix(lua: &mut hlua::Lua, _: Arc<State>) {
+    lua.set("time_unix", hlua::function0(move || -> i32 {
+        let now = Utc::now().naive_utc();
+        now.timestamp() as i32
+    }))
+}
+
 #[cfg(test)]
 mod tests {
     use crate::engine::ctx::Script;
@@ -29,4 +55,47 @@ mod tests {
         "#).expect("Failed to load script");
         script.test().expect("Script failed");
     }
+
+    #[test]
+    fn verify_strftime() {
+        let script = Script::load_unchecked(r#"
+        function run()
+            t = strftime('%d/%m/%Y %H:%M', 1558584994)
+            print(t)
+            if t ~= '23/05/2019 04:16' then
+                return 'wrong time'
+            end
+        end
+        "#).expect("Failed to load script");
+        script.test().expect("Script failed");
+
+    }
+
+    #[test]
+    fn verify_strptime() {
+        let script = Script::load_unchecked(r#"
+        function run()
+            t = strptime('%d/%m/%Y %H:%M', '23/05/2019 04:16')
+            print(t)
+            if t ~= 1558584960 then
+                return 'wrong time'
+            end
+        end
+        "#).expect("Failed to load script");
+        script.test().expect("Script failed");
+    }
+
+    #[test]
+    fn verify_time_unix() {
+        let script = Script::load_unchecked(r#"
+        function run()
+            now = time_unix()
+            print(now)
+            if now <= 1558584994 then
+                return 'time went backwards'
+            end
+        end
+        "#).expect("Failed to load script");
+        script.test().expect("Script failed");
+    }
 }