【学習メモ】Rustで作るCLIツール

お盆休み最終日！
今日は、ここまでのRust学習のまとめとして、CLIツールを作ってみたいと思います！

仕様

• AES256およびその他のアルゴリズムに対応した暗号化ツール
• コマンドライン引数から、入力ファイル名、出力ファイル名、アルゴリズムを指定できる
• 同じプログラムで暗号化と復号化の両方ができる

依存関係

以下のような依存関係になります。

Cargo.toml

[dependencies]
aes = "0.8.4"
block-padding = "0.3.3"
cbc = { version = "0.1.2", features = ["std"] }
clap = { version = "4.5.16", features = ["derive"] }

実装

暗号化・復号化部分。
複数のアルゴリズムに対応できるように、トレイトを使っています。

AESのキーや初期ベクトルは適当な値を入れています。
くれぐれも実案件ではこのコードをコピペしないでくださいねー！

encrypt.rs

use aes::cipher::{BlockDecryptMut, BlockEncryptMut, KeyIvInit};
use block_padding::{Pkcs7, UnpadError};

type Aes256Enc = cbc::Encryptor<aes::Aes256>;
type Aes256Dec = cbc::Decryptor<aes::Aes256>;

pub trait Encryptor {
    fn encrypt(&self, data: &Vec<u8>) -> Vec<u8>;
    fn decrypt(&self, data: &Vec<u8>) -> Result<Vec<u8>, UnpadError>;
}

pub struct Aes256Encryptor {
    key: Vec<u8>,
    iv: Vec<u8>
}

impl Aes256Encryptor {
    pub fn new() -> Self {
        let key = b"tekitou na password desu 1234567".to_vec();
        let iv = b"tekitou na vec 1".to_vec();
        Self {key, iv}
    }
}

impl Encryptor for Aes256Encryptor {
    fn encrypt(&self, data: &Vec<u8>) -> Vec<u8> {
        Aes256Enc::new_from_slices(&self.key, &self.iv)
            .unwrap()
            .encrypt_padded_vec_mut::<Pkcs7>(data)
    }

    fn decrypt(&self, data: &Vec<u8>) -> Result<Vec<u8>, UnpadError> {
        Aes256Dec::new_from_slices(&self.key, &self.iv)
            .unwrap()
            .decrypt_padded_vec_mut::<Pkcs7>(data)
    }
}

メイン関数。
コマンドライン引数をパースし、ひとつひとつ処理を進めています。
matchによって、エラーハンドルなどスッキリ書けた気がする！

main.rs

mod encrypt;

use clap::Parser;
use std::fs;
use crate::encrypt::{Aes256Encryptor, Encryptor};

#[derive(Parser, Debug)]
struct Args {
    /// モード（1が暗号化、2が復号化）
    #[arg(short, long, default_value_t = 1)]
    mode: u32,

    /// 入力ファイルのパス
    #[arg(short, long)]
    input: String,

    /// 出力ファイルのパス
    #[arg(short, long)]
    output: String,

    /// 暗号化アルゴリズム
    #[arg(short, long, default_value_t = String::from("AES"))]
    algorithm: String
}

fn main() {
    let args = Args::parse();

    println!("モード: {}", args.mode);
    println!("入力ファイル名: {}", args.input);
    println!("出力ファイル名: {}", args.output);
    println!("暗号化アルゴリズム: {}", args.algorithm);

    let encryptor = match args.algorithm.as_str() {
        "AES" => Aes256Encryptor::new(),
        _ => {
            println!("対応していない暗号モードです");
            return;
        }
    };

    let input_file = match fs::read(&args.input) {
        Ok(data) => data,
        Err(e) => {
            println!("入力ファイルの読み取りに失敗しました: {}", e.to_string());
            return;
        }
    };

    let crypted = match args.mode {
        1 => encryptor.encrypt(&input_file),
        2 => encryptor.decrypt(&input_file).unwrap(),
        _ => {
            println!("無効なモードです");
            return;
        }
    };

    match fs::write(&args.output, &crypted) {
        Ok(_) => (),
        Err(e) => {
            println!("出力ファイルの書き出しに失敗しました: {}", e.to_string());
            return;
        }
    }

    println!("完了しました")
}

使い方

cargo run -- -m 1 -i ./sample.txt -o ./output

暗号化したい時は -m を１に、復号したい時は2にします。
-i に入力ファイル名、-o に出力ファイル名を入れます。
アルゴリズムは -a で指定でき、デフォルトはAESです。

まとめ

Result型のエラーの部分の型など、コンパイル時にたくさんエラーが出て苦しみました...。
とっても難しい、でも楽しい言語ですね、Rust。
コンパイルさえ通れば、安全で高速な動作が保証されるので、根気よく頑張りたいところです。
それでは！