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