From 033942c2b6844568d2c4012e1c057922752267a9 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Wed, 27 Apr 2022 00:06:01 +0800
Subject: [PATCH] feat: RequestIdentities works

---
 src/cmd_sshagent.rs | 115 +++++++++++++++++++++++++++++++++++++-------
 1 file changed, 97 insertions(+), 18 deletions(-)

diff --git a/src/cmd_sshagent.rs b/src/cmd_sshagent.rs
index 5b095d0..b64c4d0 100644
--- a/src/cmd_sshagent.rs
+++ b/src/cmd_sshagent.rs
@@ -1,7 +1,98 @@
+use std::error::Error;
+use std::fs::remove_file;
+
 use clap::{App, Arg, ArgMatches, SubCommand};
 use openpgp_card::{KeyType, OpenPgp};
 use openpgp_card::crypto_data::PublicKeyMaterial;
+use openpgp_card_pcsc::PcscBackend;
 use rust_util::util_clap::{Command, CommandError};
+use rust_util::XResult;
+use ssh_agent::Agent;
+use ssh_agent::proto::{RsaPublicKey, to_bytes};
+use ssh_agent::proto::message::{self, Message, SignRequest};
+use ssh_agent::proto::public_key::PublicKey;
+
+use crate::sshutil::with_sign;
+
+struct SshAgent {
+    card: PcscBackend,
+    public_key: PublicKey,
+    comment: String,
+    ssh_string: String,
+}
+
+impl SshAgent {
+    fn new() -> XResult<Self> {
+        let mut card = crate::pgpcardutil::get_card()?;
+        let (public_key, comment, ssh_string) = {
+            let mut pgp = OpenPgp::new(&mut card);
+            let mut trans = opt_result!(pgp.transaction(), "Open card failed: {}");
+            let serial = trans.application_related_data()
+                .map(|d| d.application_id().map(|i| i.serial()))
+                .unwrap_or_else(|_| Ok(0)).unwrap_or_else(|_| 0);
+            let serial = hex::encode(serial.to_be_bytes());
+            let public_key = opt_result!(trans.public_key(KeyType::Signing), "Cannot find signing key: {}");
+            let rsa_public_key = match public_key {
+                PublicKeyMaterial::E(_) => return simple_error!("Not supports ec key"),
+                PublicKeyMaterial::R(rsa_public_key) => rsa_public_key,
+                _ => return simple_error!("Unknown key type"),
+            };
+            let e = rsa_public_key.v();
+            let n = rsa_public_key.n();
+
+            let public_key = PublicKey::Rsa(RsaPublicKey {
+                e: with_sign(e.to_vec()),
+                n: with_sign(n.to_vec()),
+            });
+            let comment = format!("pgp-card:{}", serial);
+            (public_key, comment.clone(), crate::sshutil::generate_ssh_string(e, n, &comment))
+        };
+        Ok(Self {
+            card,
+            public_key,
+            comment,
+            ssh_string,
+        })
+    }
+
+    fn handle_message(&self, request: Message) -> Result<Message, Box<dyn Error>> {
+        debugging!("Request: {:?}", request);
+        let response = match request {
+            Message::RequestIdentities => {
+                let mut identities = vec![];
+                identities.push(message::Identity {
+                    pubkey_blob: to_bytes(&self.public_key)?,
+                    comment: self.comment.clone(),
+                });
+                Ok(Message::IdentitiesAnswer(identities))
+            }
+            Message::RemoveIdentity(ref _identity) => {
+                Err(From::from(format!("Not supported message: {:?}", request)))
+            }
+            Message::AddIdentity(ref _identity) => {
+                Err(From::from(format!("Not supported message: {:?}", request)))
+            }
+            // Message::SignRequest(request) => {
+            //     let signature = to_bytes(&self.sign(&request)?)?;
+            //     Ok(Message::SignResponse(signature))
+            // }
+            _ => Err(From::from(format!("Unknown message: {:?}", request)))
+        };
+        debugging!("Response {:?}", response);
+        return response;
+    }
+}
+
+impl Agent for SshAgent {
+    type Error = ();
+
+    fn handle(&self, message: Message) -> Result<Message, ()> {
+        self.handle_message(message).or_else(|error| {
+            warning!("Error handling message - {:?}", error);
+            Ok(Message::Failure)
+        })
+    }
+}
 
 pub struct CommandImpl;
 
@@ -28,25 +119,13 @@ impl Command for CommandImpl {
         information!("Sock file: {}", sock_file);
 
         if use_pgp {
-            let mut card = crate::pgpcardutil::get_card()?;
-            let mut pgp = OpenPgp::new(&mut card);
-            let mut trans = opt_result!(pgp.transaction(), "Open card failed: {}");
-            let serial = trans.application_related_data()
-                .map(|d| d.application_id().map(|i| i.serial()))
-                .unwrap_or_else(|_| Ok(0)).unwrap_or_else(|_| 0);
-            let serial = hex::encode(serial.to_be_bytes());
-            let public_key = opt_result!(trans.public_key(KeyType::Signing), "Cannot find signing key: {}");
-            let rsa_public_key = match public_key {
-                PublicKeyMaterial::E(_) => return simple_error!("Not supports ec key"),
-                PublicKeyMaterial::R(rsa_public_key) => rsa_public_key,
-                _ => return simple_error!("Unknown key type"),
-            };
-            let e = rsa_public_key.v();
-            let n = rsa_public_key.n();
+            let ssh_agent = SshAgent::new()?;
+            information!("{}", &ssh_agent.ssh_string);
 
-            let comment = format!("pgp-card:{}", serial);
-            let ssh_string = crate::sshutil::generate_ssh_string(e, n, &comment);
-            information!("{}", ssh_string);
+            let _ = remove_file(sock_file);
+
+            information!("Start unix socket: {}", sock_file);
+            opt_result!(ssh_agent.run_unix(sock_file), "Run unix socket: {}, failed: {}", sock_file);
         }
 
         information!("card-cli ssh-agent...");