atuin_client/import/
zsh.rs

1// import old shell history!
2// automatically hoover up all that we can find
3
4use std::borrow::Cow;
5use std::path::PathBuf;
6
7use async_trait::async_trait;
8use directories::UserDirs;
9use eyre::{eyre, Result};
10use time::OffsetDateTime;
11
12use super::{get_histpath, unix_byte_lines, Importer, Loader};
13use crate::history::History;
14use crate::import::read_to_end;
15
16#[derive(Debug)]
17pub struct Zsh {
18    bytes: Vec<u8>,
19}
20
21fn default_histpath() -> Result<PathBuf> {
22    // oh-my-zsh sets HISTFILE=~/.zhistory
23    // zsh has no default value for this var, but uses ~/.zhistory.
24    // we could maybe be smarter about this in the future :)
25    let user_dirs = UserDirs::new().ok_or_else(|| eyre!("could not find user directories"))?;
26    let home_dir = user_dirs.home_dir();
27
28    let mut candidates = [".zhistory", ".zsh_history"].iter();
29    loop {
30        match candidates.next() {
31            Some(candidate) => {
32                let histpath = home_dir.join(candidate);
33                if histpath.exists() {
34                    break Ok(histpath);
35                }
36            }
37            None => {
38                break Err(eyre!(
39                    "Could not find history file. Try setting and exporting $HISTFILE"
40                ))
41            }
42        }
43    }
44}
45
46#[async_trait]
47impl Importer for Zsh {
48    const NAME: &'static str = "zsh";
49
50    async fn new() -> Result<Self> {
51        let bytes = read_to_end(get_histpath(default_histpath)?)?;
52        Ok(Self { bytes })
53    }
54
55    async fn entries(&mut self) -> Result<usize> {
56        Ok(super::count_lines(&self.bytes))
57    }
58
59    async fn load(self, h: &mut impl Loader) -> Result<()> {
60        let now = OffsetDateTime::now_utc();
61        let mut line = String::new();
62
63        let mut counter = 0;
64        for b in unix_byte_lines(&self.bytes) {
65            let s = match unmetafy(b) {
66                Some(s) => s,
67                _ => continue, // we can skip past things like invalid utf8
68            };
69
70            if let Some(s) = s.strip_suffix('\\') {
71                line.push_str(s);
72                line.push_str("\\\n");
73            } else {
74                line.push_str(&s);
75                let command = std::mem::take(&mut line);
76
77                if let Some(command) = command.strip_prefix(": ") {
78                    counter += 1;
79                    h.push(parse_extended(command, counter)).await?;
80                } else {
81                    let offset = time::Duration::seconds(counter);
82                    counter += 1;
83
84                    let imported = History::import()
85                        // preserve ordering
86                        .timestamp(now - offset)
87                        .command(command.trim_end().to_string());
88
89                    h.push(imported.build().into()).await?;
90                }
91            }
92        }
93
94        Ok(())
95    }
96}
97
98fn parse_extended(line: &str, counter: i64) -> History {
99    let (time, duration) = line.split_once(':').unwrap();
100    let (duration, command) = duration.split_once(';').unwrap();
101
102    let time = time
103        .parse::<i64>()
104        .ok()
105        .and_then(|t| OffsetDateTime::from_unix_timestamp(t).ok())
106        .unwrap_or_else(OffsetDateTime::now_utc)
107        + time::Duration::milliseconds(counter);
108
109    // use nanos, because why the hell not? we won't display them.
110    let duration = duration.parse::<i64>().map_or(-1, |t| t * 1_000_000_000);
111
112    let imported = History::import()
113        .timestamp(time)
114        .command(command.trim_end().to_string())
115        .duration(duration);
116
117    imported.build().into()
118}
119
120fn unmetafy(line: &[u8]) -> Option<Cow<str>> {
121    if line.contains(&0x83) {
122        let mut s = Vec::with_capacity(line.len());
123        let mut is_meta = false;
124        for ch in line {
125            if *ch == 0x83 {
126                is_meta = true;
127            } else if is_meta {
128                is_meta = false;
129                s.push(*ch ^ 32);
130            } else {
131                s.push(*ch)
132            }
133        }
134        String::from_utf8(s).ok().map(Cow::Owned)
135    } else {
136        std::str::from_utf8(line).ok().map(Cow::Borrowed)
137    }
138}
139
140#[cfg(test)]
141mod test {
142    use itertools::assert_equal;
143
144    use crate::import::tests::TestLoader;
145
146    use super::*;
147
148    #[test]
149    fn test_parse_extended_simple() {
150        let parsed = parse_extended("1613322469:0;cargo install atuin", 0);
151
152        assert_eq!(parsed.command, "cargo install atuin");
153        assert_eq!(parsed.duration, 0);
154        assert_eq!(
155            parsed.timestamp,
156            OffsetDateTime::from_unix_timestamp(1_613_322_469).unwrap()
157        );
158
159        let parsed = parse_extended("1613322469:10;cargo install atuin;cargo update", 0);
160
161        assert_eq!(parsed.command, "cargo install atuin;cargo update");
162        assert_eq!(parsed.duration, 10_000_000_000);
163        assert_eq!(
164            parsed.timestamp,
165            OffsetDateTime::from_unix_timestamp(1_613_322_469).unwrap()
166        );
167
168        let parsed = parse_extended("1613322469:10;cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷", 0);
169
170        assert_eq!(parsed.command, "cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷");
171        assert_eq!(parsed.duration, 10_000_000_000);
172        assert_eq!(
173            parsed.timestamp,
174            OffsetDateTime::from_unix_timestamp(1_613_322_469).unwrap()
175        );
176
177        let parsed = parse_extended("1613322469:10;cargo install \\n atuin\n", 0);
178
179        assert_eq!(parsed.command, "cargo install \\n atuin");
180        assert_eq!(parsed.duration, 10_000_000_000);
181        assert_eq!(
182            parsed.timestamp,
183            OffsetDateTime::from_unix_timestamp(1_613_322_469).unwrap()
184        );
185    }
186
187    #[tokio::test]
188    async fn test_parse_file() {
189        let bytes = r": 1613322469:0;cargo install atuin
190: 1613322469:10;cargo install atuin; \
191cargo update
192: 1613322469:10;cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷
193"
194        .as_bytes()
195        .to_owned();
196
197        let mut zsh = Zsh { bytes };
198        assert_eq!(zsh.entries().await.unwrap(), 4);
199
200        let mut loader = TestLoader::default();
201        zsh.load(&mut loader).await.unwrap();
202
203        assert_equal(
204            loader.buf.iter().map(|h| h.command.as_str()),
205            [
206                "cargo install atuin",
207                "cargo install atuin; \\\ncargo update",
208                "cargo :b̷i̶t̴r̵o̴t̴ ̵i̷s̴ ̷r̶e̵a̸l̷",
209            ],
210        );
211    }
212
213    #[tokio::test]
214    async fn test_parse_metafied() {
215        let bytes =
216            b"echo \xe4\xbd\x83\x80\xe5\xa5\xbd\nls ~/\xe9\x83\xbf\xb3\xe4\xb9\x83\xb0\n".to_vec();
217
218        let mut zsh = Zsh { bytes };
219        assert_eq!(zsh.entries().await.unwrap(), 2);
220
221        let mut loader = TestLoader::default();
222        zsh.load(&mut loader).await.unwrap();
223
224        assert_equal(
225            loader.buf.iter().map(|h| h.command.as_str()),
226            ["echo 你好", "ls ~/音乐"],
227        );
228    }
229}