From 00b7a550ab4111679af065f64d4f3e169e31e35c Mon Sep 17 00:00:00 2001
From: "xiongdian.me" <xiongdian.me@bytedance.com>
Date: Mon, 22 May 2023 19:05:24 +0800
Subject: [PATCH] riot: support zssh and sftp, and set ssh options according to
 trust settings

---
 scripts/riot | 31 +++++++++++++++++++++++++++----
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/scripts/riot b/scripts/riot
index d72ae1c..97cff0c 100755
--- a/scripts/riot
+++ b/scripts/riot
@@ -2,16 +2,22 @@
 # connect to iot services
 THIS_DIR=$( cd "$( dirname "${BASH_SOURCE[0]:-${(%):-%x}}" )" && pwd )
 source "$THIS_DIR/../tools/common.sh"
+RIOT_TRUST_CLIENT=${RIOT_TRUST_CLIENT:-${DFS_TRUST:-0}}
+RIOT_TRUST_SERVER=${RIOT_TRUST_SERVER:-0}
 
 # get target settings
 # provides:
 SERVER=""
 PORT=""  # optional
 USERNAME=""  # optional
-SSH_OPTIONS=""  # optional
+SSH_OPTIONS=""
+if [[ "$RIOT_TRUST_CLIENT" == "1" ]]; then
+    SSH_OPTIONS='-o ControlMaster=auto -o ControlPath=/tmp/sshcm-%C -o PermitLocalCommand=yes'
+fi
 get_server_meta()
 {
-    arg="$1"
+    local trust_server="$RIOT_TRUST_SERVER"
+    local arg="$1"
     # overwrite
     if [[ "$arg" == *@* ]]; then
         USERNAME=${arg%%@*}
@@ -35,12 +41,14 @@ get_server_meta()
             SERVER=$host.$domain.ink
             PORT=${PORT:-12022}
             USERNAME=${USERNAME:-root}
+            trust_server=1
             ;;
         nasp )
             SERVER=$host
             PORT=${PORT:-12022}
             USERNAME=${USERNAME:-dictxiong}
-            SSH_OPTIONS='-o ProxyJump="ssh@nasp.ob.ac.cn:36022"'
+            SSH_OPTIONS=$SSH_OPTIONS' -o ProxyJump="ssh@nasp.ob.ac.cn:36022"'
+            trust_server=1
             ;;
         proxied )
             SERVER=proxy.beardic.cn
@@ -48,17 +56,23 @@ get_server_meta()
             tmp=${tmp:0:4}
             PORT=$((10#$tmp+36000))
             USERNAME=root
+            trust_server=1
             ;;
         * )
             fmt_warning "unknown domain: $domain. will try as server name"
             SERVER="$arg"
     esac
+    if [[ "$trust_server" == "1" ]]; then
+        SSH_OPTIONS=$SSH_OPTIONS' -o ForwardX11=yes -o ForwardAgent=yes'
+    fi
 }
 
 # ssh
+SSH_BIN=ssh
+SSH_PORT_PARAM='-p'
 run_ssh()
 {
-    CMD="ssh ${PORT:+-p} $PORT $SSH_OPTIONS $USERNAME${USERNAME:+@}$SERVER"
+    CMD="$SSH_BIN ${PORT:+$SSH_PORT_PARAM} $PORT $SSH_OPTIONS $USERNAME${USERNAME:+@}$SERVER"
     fmt_note "-->" $CMD
     if [[ "$DFS_DRY_RUN" == "1" ]]; then
         echo $CMD
@@ -116,6 +130,15 @@ router()
         ssh|"" )
             run_ssh
             ;;
+        zssh )
+            SSH_BIN=zssh
+            run_ssh
+            ;;
+        sftp )
+            SSH_BIN=sftp
+            SSH_PORT_PARAM='-P'
+            run_ssh
+            ;;
         sshl )
             run_sshl "$3"
             ;;