@@ -2,7 +2,8 @@
22import Prelude hiding (catch)
33import System.Environment (getArgs, getEnv)
44import System.IO (IOMode(ReadWriteMode), hClose)
5−import System.FilePath ((</>), takeFileName)
5+import System.FilePath ((</>))
6+import System.Posix.Types (Fd)
67import Data.IORef (IORef, writeIORef, readIORef, newIORef)
78import qualified Data.ByteString.Lazy.Char8 as BSLC (pack, putStrLn)
89import Control.Exception (SomeException, catch)
@@ -63,16 +64,14 @@ import Data.Certificate.X509 (X509)
6364-- SSH
6465import Network.SSH.Client.LibSSH2
6566 ( withSSH2
66− , scpReceiveFile
67− , scpSendFile
68− , checkHost
69− , withSession
7067 , readAllChannel
68+ , retryIfNeeded
69+ , Session
70+ , Channel
7171 )
7272import Network.SSH.Client.LibSSH2.Foreign
7373 ( initialize
7474 , exit
75− , publicKeyAuthFile
7675 , channelExecute
7776 )
7877import Codec.Binary.UTF8.String (decodeString)
@@ -85,10 +84,6 @@ main = do
8584 tryConnectToAzure subscriptionId pathToCert pathToKey
8685 ["command", user, host, port, cmd] ->
8786 runCommand user host (read port) cmd
88− ["send", user, host, port, path] ->
89− sendFile user host (read port) path
90− ["receive", user, host, port, path] ->
91− receiveFile user host (read port) path
9287 _ ->
9388 putStrLn "Invalid command line arguments"
9489
@@ -96,52 +91,24 @@ main = do
9691-- Taken from libssh2/ssh-client --
9792--------------------------------------------------------------------------------
9893
94+runCommand :: String -> String -> Int -> String -> IO ()
9995runCommand login host port command =
100− ssh login host port $ \ch -> do
101− channelExecute ch command
102− result <- readAllChannel ch
96+ ssh login host port $ \fd s ch -> do
97+ _ <- retryIfNeeded fd s $ channelExecute ch command
98+ result <- readAllChannel fd ch
10399 let r = decodeString result
104100 print (length result)
105101 print (length r)
106102 putStrLn r
107103
108−sendFile login host port path = do
109− initialize True
110− home <- getEnv "HOME"
111− let known_hosts = home </> ".ssh" </> "known_hosts"
112− public = home </> ".ssh" </> "id_rsa.pub"
113− private = home </> ".ssh" </> "id_rsa"
114−
115− withSession host port $ \_ s -> do
116− r <- checkHost s host port known_hosts
117− print r
118− publicKeyAuthFile s login public private ""
119− sz <- scpSendFile s 0o644 path (takeFileName path)
120− putStrLn $ "Sent: " ++ show sz ++ " bytes."
121− exit
122−
123−receiveFile login host port path = do
124− initialize True
125− home <- getEnv "HOME"
126− let known_hosts = home </> ".ssh" </> "known_hosts"
127− public = home </> ".ssh" </> "id_rsa.pub"
128− private = home </> ".ssh" </> "id_rsa"
129−
130− withSession host port $ \_ s -> do
131− r <- checkHost s host port known_hosts
132− print r
133− publicKeyAuthFile s login public private ""
134− sz <- scpReceiveFile s (takeFileName path) path
135− putStrLn $ "Received: " ++ show sz ++ " bytes."
136− exit
137−
104+ssh :: String -> String -> Int -> (Fd -> Session -> Channel -> IO a) -> IO ()
138105ssh login host port actions = do
139− initialize True
106+ _ <- initialize True
140107 home <- getEnv "HOME"
141108 let known_hosts = home </> ".ssh" </> "known_hosts"
142109 public = home </> ".ssh" </> "id_rsa.pub"
143110 private = home </> ".ssh" </> "id_rsa"
144− withSSH2 known_hosts public private login host port $ actions
111+ _ <- withSSH2 known_hosts public private login host port $ actions
145112 exit
146113
147114--------------------------------------------------------------------------------