1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

// Copyright 2020 - developers of the `grammers` project.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
use grammers_crypto::auth_key::AuthKey;
use std::fs::{File, OpenOptions};
use std::io::{self, BufRead, BufReader, Seek, Write};
use std::path::Path;

const CURRENT_VERSION: u32 = 1;

fn parse_hex(byte: &str) -> Option<u8> {
    match u8::from_str_radix(byte, 16) {
        Ok(x) => Some(x),
        Err(_) => None,
    }
}

fn key_from_hex(hex: &str) -> Option<[u8; 256]> {
    let mut buffer = [0; 256];
    if hex.len() == buffer.len() * 2 {
        for (i, byte) in buffer.iter_mut().enumerate() {
            let i = i * 2;
            if let Some(value) = parse_hex(&hex[i..i + 2]) {
                *byte = value;
            } else {
                return None;
            }
        }

        Some(buffer)
    } else {
        None
    }
}

fn hex_from_key(key: &[u8; 256]) -> String {
    use std::fmt::Write;
    let mut buffer = String::with_capacity(key.len() * 2);
    for byte in key.iter() {
        write!(buffer, "{:02x}", byte).unwrap();
    }
    buffer
}

pub struct Session {
    file: File,
    pub user_dc: Option<i32>,
    pub auth_key: Option<AuthKey>,
}

impl Session {
    /// Loads or creates a new session file.
    pub fn load_or_create<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        if let Ok(instance) = Self::load(path.as_ref()) {
            Ok(instance)
        } else {
            Self::create(path)
        }
    }

    /// Create a new session instance.
    pub fn create<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        Ok(Self {
            file: File::create(path)?,
            user_dc: None,
            auth_key: None,
        })
    }

    /// Load a previous session instance.
    pub fn load<P: AsRef<Path>>(path: P) -> io::Result<Self> {
        let mut lines = BufReader::new(File::open(path.as_ref())?).lines();

        // Version
        let version: u32 = if let Some(Ok(line)) = lines.next() {
            line.parse()
                .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "malformed session"))?
        } else {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "malformed session",
            ));
        };
        if version != CURRENT_VERSION {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "unknown version",
            ));
        }

        // user_dc
        let user_dc = if let Some(Ok(line)) = lines.next() {
            match line.parse() {
                Ok(x) => Some(x),
                Err(_) => None,
            }
        } else {
            None
        };

        // auth_key
        let auth_key = if let Some(Ok(line)) = lines.next() {
            key_from_hex(&line)
        } else {
            None
        };

        drop(lines);
        Ok(Self {
            file: OpenOptions::new().write(true).open(path.as_ref())?,
            user_dc,
            auth_key: auth_key.map(AuthKey::from_bytes),
        })
    }

    /// Saves the session file.
    pub fn save(&mut self) -> io::Result<()> {
        self.file.seek(io::SeekFrom::Start(0))?;
        writeln!(self.file, "{}", CURRENT_VERSION)?;

        if let Some(dc_id) = self.user_dc {
            writeln!(self.file, "{}", dc_id)?;
        } else {
            writeln!(self.file)?;
        }

        if let Some(data) = &self.auth_key {
            writeln!(self.file, "{}", hex_from_key(&data.to_bytes()))?;
        } else {
            writeln!(self.file)?;
        }
        self.file.sync_data()?;
        Ok(())
    }
}