ANSIColour.hs

   1 -- This file is part of Diohsc
   2 -- Copyright (C) 2020 Martin Bays <mbays@sdf.org>
   3 --
   4 -- This program is free software: you can redistribute it and/or modify
   5 -- it under the terms of version 3 of the GNU General Public License as
   6 -- published by the Free Software Foundation, or any later version.
   7 --
   8 -- You should have received a copy of the GNU General Public License
   9 -- along with this program.  If not, see http://www.gnu.org/licenses/.
  10
  11 {-# LANGUAGE BangPatterns      #-}
  12 {-# LANGUAGE LambdaCase        #-}
  13 {-# LANGUAGE OverloadedStrings #-}
  14 {-# LANGUAGE Safe              #-}
  15
  16 -- Basic ansi attributes, using only most widely supported ansi terminal codes
  17 module ANSIColour
  18     ( applyIf
  19     , resetCode
  20     , withColour
  21     , withBold
  22     , withReverse
  23     , withUnderline
  24     , withColourStr
  25     , withBoldStr
  26     , withReverseStr
  27     , withUnderlineStr
  28     , stripCSI
  29     , visibleLength
  30     , splitAtVisible
  31     , escapePromptCSI
  32     , sanitiseForDisplay
  33     , stripControl
  34     , stripControlExceptTab
  35     , Colour(..)
  36     ) where
  37
  38 import           Control.Exception.Base (bracket_)
  39
  40 import qualified Data.Text.Lazy         as T
  41 import qualified Data.Text.Lazy.IO      as T
  42
  43 import           MetaString
  44 import           WCWidth
  45
  46 data Colour = Black | Red | Green | Yellow
  47     | Blue | Magenta | Cyan | White
  48     | BoldBlack | BoldRed | BoldGreen | BoldYellow
  49     | BoldBlue | BoldMagenta | BoldCyan | BoldWhite
  50     deriving (Eq,Ord,Show,Read)
  51
  52 resetCode, boldCode, unboldCode, reverseCode,
  53     unreverseCode, underlineCode , ununderlineCode, resetColourCode
  54     :: MetaString a => a
  55 resetCode = "\ESC[0m"
  56 boldCode = "\ESC[1m"
  57 underlineCode = "\ESC[4m"
  58 reverseCode = "\ESC[7m"
  59 unboldCode = "\ESC[22m"
  60 ununderlineCode = "\ESC[24m"
  61 unreverseCode = "\ESC[27m"
  62 resetColourCode = "\ESC[39m\ESC[22m"
  63
  64 colourCode :: MetaString a => Colour -> a
  65 colourCode c = (if isBold c then boldCode else "") <> "\ESC[3" <> fromString (colNum c) <> "m"
  66     where
  67     isBold = flip elem [BoldBlack, BoldRed, BoldGreen, BoldYellow,
  68         BoldBlue, BoldMagenta, BoldCyan, BoldWhite]
  69     colNum Black       = "0"
  70     colNum Red         = "1"
  71     colNum Green       = "2"
  72     colNum Yellow      = "3"
  73     colNum Blue        = "4"
  74     colNum Magenta     = "5"
  75     colNum Cyan        = "6"
  76     colNum White       = "7"
  77     colNum BoldBlack   = "0"
  78     colNum BoldRed     = "1"
  79     colNum BoldGreen   = "2"
  80     colNum BoldYellow  = "3"
  81     colNum BoldBlue    = "4"
  82     colNum BoldMagenta = "5"
  83     colNum BoldCyan    = "6"
  84     colNum BoldWhite   = "7"
  85
  86 withStyle :: T.Text -> T.Text -> IO a -> IO a
  87 withStyle c r = T.putStr c `bracket_` T.putStr r
  88 withColour :: Colour -> IO a -> IO a
  89 withColour c = withStyle (colourCode c) resetColourCode
  90 withBold, withReverse, withUnderline :: IO a -> IO a
  91 withBold = withStyle boldCode unboldCode
  92 withReverse = withStyle reverseCode unreverseCode
  93 withUnderline = withStyle underlineCode ununderlineCode
  94
  95 withStyleStr :: MetaString a => a -> a -> a -> a
  96 withStyleStr c r s = c <> s <> r
  97 withColourStr :: MetaString a => Colour -> a -> a
  98 withColourStr c = withStyleStr (colourCode c) resetColourCode
  99 withBoldStr, withReverseStr, withUnderlineStr :: MetaString a => a -> a
 100 withBoldStr = withStyleStr boldCode unboldCode
 101 withReverseStr = withStyleStr reverseCode unreverseCode
 102 withUnderlineStr = withStyleStr underlineCode ununderlineCode
 103
 104 -- |"applyIf cond f" is shorthand for "if cond then f else id"
 105 applyIf :: Bool -> (a -> a) -> (a -> a)
 106 applyIf True  = id
 107 applyIf False = const id
 108
 109 mapTail :: (a -> a) -> [a] -> [a]
 110 mapTail _ []     = []
 111 mapTail f (a:as) = a:(f<$>as)
 112
 113
 114 endCSI :: Char -> Bool
 115 endCSI c = '@' <= c && c <= '~'
 116
 117 -- |strip all CSI escape sequences
 118 stripCSI :: T.Text -> T.Text
 119 stripCSI = T.concat . mapTail dropCSI . T.splitOn "\ESC["
 120     where dropCSI = T.drop 1 . T.dropWhile (not . endCSI)
 121
 122 visibleLength :: (Integral i) => T.Text -> i
 123 visibleLength = fromIntegral . wcLength . stripCSI
 124
 125 wcLength :: T.Text -> Int
 126 wcLength = sum . (max 0 . wcwidth <$>) . T.unpack
 127
 128 splitAtWC :: Int -> T.Text -> (T.Text,T.Text)
 129 splitAtWC m = go m T.empty where
 130     go !n !acc t
 131         | Just (c,r) <- T.uncons t = let w = max 0 $ wcwidth c in
 132             if w > max 0 n then (T.reverse acc,t)
 133             else go (n - w) (T.cons c acc) r
 134         | otherwise = (T.reverse acc, T.empty)
 135
 136
 137 splitAtVisible :: (Integral i) => i -> T.Text -> (T.Text,T.Text)
 138 splitAtVisible n t =
 139     let (pre,post) = T.breakOn "\ESC[" t
 140         n' = fromIntegral n
 141         (a,b) = splitAtWC n' pre
 142         catFst s (s',s'') = (s<>s',s'')
 143     in a `catFst`
 144         if not (T.null b) || T.null post then ("",b<>post)
 145         else
 146             let (s,r) = T.splitAt 2 post
 147                 (s',r') = T.break endCSI r
 148                 (s'',rest) = T.splitAt 1 r'
 149                 csi = s <> s' <> s''
 150             in csi `catFst` splitAtVisible (n' - wcLength a) rest
 151
 152 -- |sanitise non-CSI escape sequences by turning \ESC into \\ESC
 153 -- (buggy terminals make these sequences a potential security hole;
 154 -- see e.g. https://nvd.nist.gov/vuln/detail/CVE-2020-9366 )
 155 sanitiseNonCSI :: T.Text -> T.Text
 156 sanitiseNonCSI = T.concat . mapTail escapeEsc . T.splitOn "\ESC"
 157     where
 158     escapeEsc s = case T.take 1 s of
 159         "[" -> T.cons '\ESC' s
 160         _   -> "\\ESC" <> s
 161
 162 -- |strip all C0 and C1 control chars, replacing tab with space
 163 stripControl :: T.Text -> T.Text
 164 stripControl = T.concatMap $ \case
 165     '\t'               -> " "
 166     c | wcwidth c == 0 -> ""
 167     c                  -> T.singleton c
 168
 169 -- |strip all C0 and C1 control chars except '\t'
 170 stripControlExceptTab :: T.Text -> T.Text
 171 stripControlExceptTab = T.filter $ \c -> c == '\t' || wcwidth c > 0
 172
 173 -- |strip all C0 and C1 control chars except tab, and esc where it introduces
 174 -- a CSI escape sequence. (Might be even better to strip all but SGR, but that
 175 -- would require more parsing.)
 176 sanitiseForDisplay :: T.Text -> T.Text
 177 sanitiseForDisplay = sanitiseNonCSI . T.filter (\c -> c `elem` ['\ESC','\t'] || wcwidth c >= 0)
 178
 179 -- |append \STX to each CSI sequence, as required in Haskeline prompts.
 180 -- See https://github.com/judah/haskeline/wiki/ControlSequencesInPrompt
 181 escapePromptCSI :: String -> String
 182 escapePromptCSI s = case break (== '\ESC') s of
 183     (pre,'\ESC':'[':post) -> ((pre <> "\ESC[") <>) $
 184         case break endCSI post of
 185             (pre',e:post') -> (pre' <>) $ e : '\STX' : escapePromptCSI post'
 186             (pre',[])      -> pre'
 187     (pre,[]) -> pre
 188     (pre,e:post) -> (pre <>) $ e : escapePromptCSI post