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