1 files changed, 15 insertions, 0 deletions

diff --git a/lib/SSHKey.hs b/lib/SSHKey.hs
index 488f55f..bd47169 100644
--- a/lib/SSHKey.hs
+++ b/lib/SSHKey.hs
@@ -1,9 +1,16 @@
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE CPP #-}
 module SSHKey where
 
 import qualified Data.ByteString.Lazy.Char8 as L8
 import qualified Data.ByteString.Lazy as L
+#if defined(VERSION_memory)
+import qualified Data.ByteString.Char8 as S8
+import Data.ByteArray.Encoding
+import FunctorToMaybe
+#elif defined(VERSION_dataenc)
 import qualified Codec.Binary.Base64 as Base64
+#endif
 import Data.Binary.Get ( runGet )
 import Data.Binary.Put ( putWord32be, runPut, putByteString )
 import Data.Binary     ( get, put )
@@ -19,7 +26,11 @@ keyblob :: Key -> L.ByteString
 keyblob (n,e) = "ssh-rsa " <> blob
  where
     bs = sshrsa e n
+#if defined(VERSION_memory)
+    blob = L.fromStrict $ convertToBase Base64 (L.toStrict bs)
+#elif defined(VERSION_dataenc)
     blob = L8.pack $ Base64.encode (L.unpack bs)
+#endif
 
     sshrsa :: Integer -> Integer -> L.ByteString
     sshrsa e n = runPut $ do
@@ -35,7 +46,11 @@ blobkey bs = do
     let (sp,bs2) = L8.span isSpace bs1
     guard $ not (L8.null sp)
     bs3 <- listToMaybe $ L8.words bs2
+#if defined(VERSION_memory)
+    qq <- fmap L.fromStrict $ functorToMaybe $ convertFromBase Base64 $ L.toStrict bs3
+#elif defined(VERSION_dataenc)
     qq <- L.pack `fmap` Base64.decode (L8.unpack bs3)
+#endif
     decode_sshrsa qq
  where
     decode_sshrsa :: L8.ByteString -> Maybe Key