module Net.TFTP_Client(tftpGet,tftpPut) where

-- TFTP Client, protocol described in RFC 1350
-- See http://www.networksorcery.com/enp/rfc/rfc1350.txt

import Net.Concurrent

import Net.UDP_Client as UDP(Packet(..),template,listenAny,unlisten)
import qualified Net.PortNumber as Port
import Net.TFTP
import qualified Net.Interface as Net
import Net.Packet(listArray)
import Net.PacketParsing(doUnparse,doParse)
import Net.Utils(arraySize)
--import Monad.Util(whileM)

tftpGet :: p
-> Interface m
-> Addr
-> [Char]
-> [Char]
-> m (Either [Char] [Data])
tftpGet p
debug Interface m
udp Addr
serverIP [Char]
filename [Char]
mode =
  do (TimedInterface m Packet Packet
tftp,m ()
close) <- Interface m -> Addr -> m (TimedInterface m Packet Packet, m ())
forall {i} {o} {m :: * -> *} {v :: * -> *}.
(Parse i, Unparse o, MVarIO v m) =>
Interface m -> Addr -> m (TimedInterface m i o, m ())
initialize Interface m
udp Addr
serverIP
     Either [Char] [Data]
result <- TimedInterface m Packet Packet
-> [Data] -> BlockNr -> m (Either [Char] [Data])
forall {m :: * -> *}.
Monad m =>
TimedInterface m Packet Packet
-> [Data] -> BlockNr -> m (Either [Char] [Data])
getBlocks TimedInterface m Packet Packet
tftp [] BlockNr
0
     m ()
close
     Either [Char] [Data] -> m (Either [Char] [Data])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Either [Char] [Data]
result
  where
    rrq :: Packet
rrq = [Char] -> [Char] -> Packet
RRQ [Char]
filename [Char]
mode

    getBlocks :: TimedInterface m Packet Packet
-> [Data] -> BlockNr -> m (Either [Char] [Data])
getBlocks TimedInterface m Packet Packet
tftp [Data]
bs BlockNr
last =
      do Maybe Packet
p <- TimedInterface m Packet Packet -> Packet -> m (Maybe Packet)
forall {m :: * -> *} {i} {o}.
Monad m =>
TimedInterface m i o -> o -> m (Maybe i)
txRetry TimedInterface m Packet Packet
tftp (if BlockNr
lastBlockNr -> BlockNr -> Bool
forall a. Eq a => a -> a -> Bool
==BlockNr
0 then Packet
rrq else BlockNr -> Packet
Ack BlockNr
last)
	 case Maybe Packet
p of
	   Maybe Packet
Nothing -> Either [Char] [Data] -> m (Either [Char] [Data])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] [Data] -> m (Either [Char] [Data]))
-> Either [Char] [Data] -> m (Either [Char] [Data])
forall a b. (a -> b) -> a -> b
$ [Char] -> Either [Char] [Data]
forall a b. a -> Either a b
Left [Char]
"Timeout"
           Just (Data BlockNr
nr Data
b) ->
	       if BlockNr
nrBlockNr -> BlockNr -> Bool
forall a. Eq a => a -> a -> Bool
==BlockNr
lastBlockNr -> BlockNr -> BlockNr
forall a. Num a => a -> a -> a
+BlockNr
1
	       then if Data -> Int
forall {a1} {a2 :: * -> * -> *} {e}.
(Num a1, IArray a2 e, Ix a1) =>
a2 a1 e -> a1
arraySize Data
bInt -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
512
		    then do TimedInterface m Packet Packet -> Packet -> m ()
forall (m :: * -> *) i o. TimedInterface m i o -> o -> m ()
Net.txT TimedInterface m Packet Packet
tftp (BlockNr -> Packet
Ack BlockNr
nr)
			    Either [Char] [Data] -> m (Either [Char] [Data])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Data] -> Either [Char] [Data]
forall a b. b -> Either a b
Right ([Data] -> [Data]
forall a. [a] -> [a]
reverse (Data
bData -> [Data] -> [Data]
forall a. a -> [a] -> [a]
:[Data]
bs)))
		    else TimedInterface m Packet Packet
-> [Data] -> BlockNr -> m (Either [Char] [Data])
getBlocks TimedInterface m Packet Packet
tftp (Data
bData -> [Data] -> [Data]
forall a. a -> [a] -> [a]
:[Data]
bs) BlockNr
nr
	       else if BlockNr
nrBlockNr -> BlockNr -> Bool
forall a. Eq a => a -> a -> Bool
==BlockNr
last
		    then TimedInterface m Packet Packet
-> [Data] -> BlockNr -> m (Either [Char] [Data])
getBlocks TimedInterface m Packet Packet
tftp [Data]
bs BlockNr
last -- ignore dupl
		    else Either [Char] [Data] -> m (Either [Char] [Data])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> Either [Char] [Data]
forall a b. a -> Either a b
Left [Char]
"unexpected data block")
	   Just (Error ErrorCode
c [Char]
msg) ->
	       Either [Char] [Data] -> m (Either [Char] [Data])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] [Data] -> m (Either [Char] [Data]))
-> Either [Char] [Data] -> m (Either [Char] [Data])
forall a b. (a -> b) -> a -> b
$ [Char] -> Either [Char] [Data]
forall a b. a -> Either a b
Left ([Char] -> Either [Char] [Data]) -> [Char] -> Either [Char] [Data]
forall a b. (a -> b) -> a -> b
$ [Char]
"Server said: "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ErrorCode -> [Char]
forall a. Show a => a -> [Char]
show ErrorCode
c[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
" "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
msg
	   Just Packet
msg -> Either [Char] [Data] -> m (Either [Char] [Data])
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] [Data] -> m (Either [Char] [Data]))
-> Either [Char] [Data] -> m (Either [Char] [Data])
forall a b. (a -> b) -> a -> b
$ [Char] -> Either [Char] [Data]
forall a b. a -> Either a b
Left ([Char] -> Either [Char] [Data]) -> [Char] -> Either [Char] [Data]
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected packet: "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++Packet -> [Char]
forall a. Show a => a -> [Char]
show Packet
msg

tftpPut :: p
-> Interface m
-> Addr
-> [Char]
-> [Char]
-> [Char]
-> m (Either [Char] ())
tftpPut p
debug Interface m
udp Addr
serverIP [Char]
filename [Char]
mode [Char]
contents =
  do (TimedInterface m Packet Packet
tftp,m ()
close) <- Interface m -> Addr -> m (TimedInterface m Packet Packet, m ())
forall {i} {o} {m :: * -> *} {v :: * -> *}.
(Parse i, Unparse o, MVarIO v m) =>
Interface m -> Addr -> m (TimedInterface m i o, m ())
initialize Interface m
udp Addr
serverIP
     Either [Char] ()
result <- TimedInterface m Packet Packet
-> [Data] -> Packet -> BlockNr -> m (Either [Char] ())
forall {m :: * -> *}.
Monad m =>
TimedInterface m Packet Packet
-> [Data] -> Packet -> BlockNr -> m (Either [Char] ())
putBlocks TimedInterface m Packet Packet
tftp ([Char] -> [Data]
blocks [Char]
contents) ([Char] -> [Char] -> Packet
WRQ [Char]
filename [Char]
mode) BlockNr
0
     m ()
close
     Either [Char] () -> m (Either [Char] ())
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Either [Char] ()
result
  where
    putBlocks :: TimedInterface m Packet Packet
-> [Data] -> Packet -> BlockNr -> m (Either [Char] ())
putBlocks TimedInterface m Packet Packet
tftp [Data]
bs Packet
packet BlockNr
current =
      do --debug $ "tftpPut send "++show packet
         Maybe Packet
p <- TimedInterface m Packet Packet -> Packet -> m (Maybe Packet)
forall {m :: * -> *} {i} {o}.
Monad m =>
TimedInterface m i o -> o -> m (Maybe i)
txRetry TimedInterface m Packet Packet
tftp Packet
packet
         --debug $ "tftpPut receive "++show p
	 case Maybe Packet
p of
	   Maybe Packet
Nothing -> Either [Char] () -> m (Either [Char] ())
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] () -> m (Either [Char] ()))
-> Either [Char] () -> m (Either [Char] ())
forall a b. (a -> b) -> a -> b
$ [Char] -> Either [Char] ()
forall a b. a -> Either a b
Left [Char]
"Timeout"
           Just (Ack BlockNr
nr) ->
	       if BlockNr
nrBlockNr -> BlockNr -> Bool
forall a. Eq a => a -> a -> Bool
==BlockNr
current
	       then case [Data]
bs of
                      [] -> Either [Char] () -> m (Either [Char] ())
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (() -> Either [Char] ()
forall a b. b -> Either a b
Right ())
                      Data
b:[Data]
bs -> TimedInterface m Packet Packet
-> [Data] -> Packet -> BlockNr -> m (Either [Char] ())
putBlocks TimedInterface m Packet Packet
tftp [Data]
bs (BlockNr -> Data -> Packet
Data BlockNr
nr Data
b) BlockNr
nr
		        where nr :: BlockNr
nr=BlockNr
currentBlockNr -> BlockNr -> BlockNr
forall a. Num a => a -> a -> a
+BlockNr
1
	       else if BlockNr
nrBlockNr -> BlockNr -> Bool
forall a. Eq a => a -> a -> Bool
==BlockNr
currentBlockNr -> BlockNr -> BlockNr
forall a. Num a => a -> a -> a
-BlockNr
1
		    then TimedInterface m Packet Packet
-> [Data] -> Packet -> BlockNr -> m (Either [Char] ())
putBlocks TimedInterface m Packet Packet
tftp [Data]
bs Packet
packet BlockNr
current -- hmm, ignore dupl
		    else Either [Char] () -> m (Either [Char] ())
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ([Char] -> Either [Char] ()
forall a b. a -> Either a b
Left [Char]
"unexpected ack")
	   Just (Error ErrorCode
c [Char]
msg) ->
	       Either [Char] () -> m (Either [Char] ())
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] () -> m (Either [Char] ()))
-> Either [Char] () -> m (Either [Char] ())
forall a b. (a -> b) -> a -> b
$ [Char] -> Either [Char] ()
forall a b. a -> Either a b
Left ([Char] -> Either [Char] ()) -> [Char] -> Either [Char] ()
forall a b. (a -> b) -> a -> b
$ [Char]
"Server said: "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ErrorCode -> [Char]
forall a. Show a => a -> [Char]
show ErrorCode
c[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
" "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++[Char]
msg
	   Just Packet
msg -> Either [Char] () -> m (Either [Char] ())
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Either [Char] () -> m (Either [Char] ()))
-> Either [Char] () -> m (Either [Char] ())
forall a b. (a -> b) -> a -> b
$ [Char] -> Either [Char] ()
forall a b. a -> Either a b
Left ([Char] -> Either [Char] ()) -> [Char] -> Either [Char] ()
forall a b. (a -> b) -> a -> b
$ [Char]
"Unexpected packet: "[Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++Packet -> [Char]
forall a. Show a => a -> [Char]
show Packet
msg

    blocks :: String -> [Data]
    blocks :: [Char] -> [Data]
blocks [Char]
s =
      case Int -> [Char] -> ([Char], [Char])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
512 [Char]
s of
	([],[Char]
_) -> []
	([Char]
s1,[Char]
s2) -> (Int, Int) -> [Word8] -> Data
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
listArray (Int
1,[Char] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Char]
s1) ([Char] -> [Word8]
conv [Char]
s1)Data -> [Data] -> [Data]
forall a. a -> [a] -> [a]
:[Char] -> [Data]
blocks [Char]
s2
      where
        conv :: [Char] -> [Word8]
conv = (Char -> Word8) -> [Char] -> [Word8]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral(Int -> Word8) -> (Char -> Int) -> Char -> Word8
forall b c a. (b -> c) -> (a -> b) -> a -> c
.Char -> Int
forall a. Enum a => a -> Int
fromEnum)

txRetry :: TimedInterface m i o -> o -> m (Maybe i)
txRetry TimedInterface m i o
tftp o
packet = Integer -> m (Maybe i)
forall {t}. (Ord t, Num t) => t -> m (Maybe i)
rx Integer
2
  where
     rx :: t -> m (Maybe i)
rx t
n =
       do TimedInterface m i o -> o -> m ()
forall (m :: * -> *) i o. TimedInterface m i o -> o -> m ()
Net.txT TimedInterface m i o
tftp o
packet
          m (Maybe i) -> (i -> m (Maybe i)) -> Maybe i -> m (Maybe i)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe m (Maybe i)
retry (Maybe i -> m (Maybe i)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe i -> m (Maybe i)) -> (i -> Maybe i) -> i -> m (Maybe i)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. i -> Maybe i
forall a. a -> Maybe a
Just) (Maybe i -> m (Maybe i)) -> m (Maybe i) -> m (Maybe i)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TimedInterface m i o -> Maybe Int -> m (Maybe i)
forall (m :: * -> *) i o.
TimedInterface m i o -> Maybe Int -> m (Maybe i)
Net.rxT TimedInterface m i o
tftp (Int -> Maybe Int
forall a. a -> Maybe a
Just Int
t)
       where
         retry :: m (Maybe i)
retry = if t
nt -> t -> Bool
forall a. Ord a => a -> a -> Bool
>t
0 then t -> m (Maybe i)
rx (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1) else Maybe i -> m (Maybe i)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe i
forall a. Maybe a
Nothing
         t :: Int
t = Int
2000000 -- microseconds

{-
txRetry tftp packet = rx 2
  where
     rx n =
	 do waiting <- newRef True
	    fork $ whileM (readRef waiting) $ Net.tx tftp packet >> delay t
	    p <- Net.rx tftp
	    writeRef waiting False
	    return (Just p)
       where
         t = 2000000 -- microseconds
-}

initialize :: Interface m -> Addr -> m (TimedInterface m i o, m ())
initialize Interface m
udp Addr
serverIP =
  do (Port
port,UDP_API m
uclient) <- Interface m -> m (Port, UDP_API m)
forall (m :: * -> *). Interface m -> m (Port, UDP_API m)
UDP.listenAny Interface m
udp
     v (Maybe Port)
portM <- Maybe Port -> m (v (Maybe Port))
forall a. a -> m (v a)
forall (v :: * -> *) (io :: * -> *) a. MVarIO v io => a -> io (v a)
newMVar Maybe Port
forall a. Maybe a
Nothing
     let rx :: Maybe Int -> m (Maybe a)
rx Maybe Int
t =
	   do Maybe (Addr, Packet InPacket)
r <- UDP_API m -> Maybe Int -> m (Maybe (Addr, Packet InPacket))
forall (m :: * -> *) i o.
TimedInterface m i o -> Maybe Int -> m (Maybe i)
Net.rxT UDP_API m
uclient Maybe Int
t
	      case Maybe (Addr, Packet InPacket)
r of
		Maybe (Addr, Packet InPacket)
Nothing -> Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
		Just (Addr
fromIP,Packet InPacket
udpP) ->
		  let sPort :: Port
sPort = Packet InPacket -> Port
forall content. Packet content -> Port
sourcePort Packet InPacket
udpP in
		  case InPacket -> Maybe a
forall {a}. Parse a => InPacket -> Maybe a
doParse (Packet InPacket -> InPacket
forall content. Packet content -> content
content Packet InPacket
udpP) of
		    Maybe a
Nothing -> Maybe Int -> m (Maybe a)
rx Maybe Int
t -- reduce t!!
		    Just a
p -> if Addr
fromIPAddr -> Addr -> Bool
forall a. Eq a => a -> a -> Bool
==Addr
serverIP -- also check port number
			      then do Maybe Port
optsp <- v (Maybe Port) -> m (Maybe Port)
forall a. v a -> m a
forall (v :: * -> *) (io :: * -> *) a. MVarIO v io => v a -> io a
takeMVar v (Maybe Port)
portM
				      case Maybe Port
optsp of
					Maybe Port
Nothing ->
					    do v (Maybe Port) -> Maybe Port -> m ()
forall a. v a -> a -> m ()
forall (v :: * -> *) (io :: * -> *) a.
MVarIO v io =>
v a -> a -> io ()
putMVar v (Maybe Port)
portM (Port -> Maybe Port
forall a. a -> Maybe a
Just Port
sPort)
					       Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
p)
					Just Port
port ->
					    do v (Maybe Port) -> Maybe Port -> m ()
forall a. v a -> a -> m ()
forall (v :: * -> *) (io :: * -> *) a.
MVarIO v io =>
v a -> a -> io ()
putMVar v (Maybe Port)
portM Maybe Port
optsp
					       if Port
portPort -> Port -> Bool
forall a. Eq a => a -> a -> Bool
==Port
sPort
						  then Maybe a -> m (Maybe a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> Maybe a
forall a. a -> Maybe a
Just a
p)
						  else Maybe Int -> m (Maybe a)
rx Maybe Int
t -- wrong port!
			      else Maybe Int -> m (Maybe a)
rx Maybe Int
t
	 tx :: p -> m ()
tx p
msg =
	     -- The initial request is sent to serverPort. After that,
	     -- messages are sent to the port chosen by the server.
	     do Port
sPort <- Port -> (Port -> Port) -> Maybe Port -> Port
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Port
serverPort Port -> Port
forall a. a -> a
id (Maybe Port -> Port) -> m (Maybe Port) -> m Port
forall a b. (a -> b) -> m a -> m b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
`fmap` v (Maybe Port) -> m (Maybe Port)
forall a. v a -> m a
forall (v :: * -> *) (io :: * -> *) a. MVarIO v io => v a -> io a
readMVar v (Maybe Port)
portM
		UDP_API m -> (Addr, Packet OutPacket) -> m ()
forall (m :: * -> *) i o. TimedInterface m i o -> o -> m ()
Net.txT UDP_API m
uclient (Addr
serverIP,Port -> Packet OutPacket
udpP Port
sPort)
            where
              udpP :: Port -> Packet OutPacket
udpP Port
sPort = Port -> Port -> OutPacket -> Packet OutPacket
forall {content}. Port -> Port -> content -> Packet content
UDP.template Port
port Port
sPort OutPacket
bs
              bs :: OutPacket
bs = p -> OutPacket
forall {a}. Unparse a => a -> OutPacket
doUnparse p
msg

     (TimedInterface m i o, m ()) -> m (TimedInterface m i o, m ())
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return ((Maybe Int -> m (Maybe i)) -> (o -> m ()) -> TimedInterface m i o
forall (m :: * -> *) i o.
(Maybe Int -> m (Maybe i)) -> (o -> m ()) -> TimedInterface m i o
Net.TimedInterface Maybe Int -> m (Maybe i)
forall {a}. Parse a => Maybe Int -> m (Maybe a)
rx o -> m ()
forall {p}. Unparse p => p -> m ()
tx,Interface m -> Port -> m ()
forall (m :: * -> *). Interface m -> Port -> m ()
unlisten Interface m
udp Port
port)

serverPort :: Port
serverPort = Port
Port.tftp -- standard TFTP server port