1# -*- coding: utf-8 -*- #
2# Copyright 2015 Google LLC. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#    http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""OS specific console_attr helper functions."""
17
18from __future__ import absolute_import
19from __future__ import division
20from __future__ import unicode_literals
21
22import os
23import sys
24
25from googlecloudsdk.core.util import encoding
26from googlecloudsdk.core.util import platforms
27
28
29def GetTermSize():
30  """Gets the terminal x and y dimensions in characters.
31
32  _GetTermSize*() helper functions taken from:
33    http://stackoverflow.com/questions/263890/
34
35  Returns:
36    (columns, lines): A tuple containing the terminal x and y dimensions.
37  """
38  xy = None
39  # Believe the first helper that doesn't bail.
40  for get_terminal_size in (_GetTermSizePosix,
41                            _GetTermSizeWindows,
42                            _GetTermSizeEnvironment,
43                            _GetTermSizeTput):
44    try:
45      xy = get_terminal_size()
46      if xy:
47        break
48    except:  # pylint: disable=bare-except
49      pass
50  return xy or (80, 24)
51
52
53def _GetTermSizePosix():
54  """Returns the Posix terminal x and y dimensions."""
55  # pylint: disable=g-import-not-at-top
56  import fcntl
57  # pylint: disable=g-import-not-at-top
58  import struct
59  # pylint: disable=g-import-not-at-top
60  import termios
61
62  def _GetXY(fd):
63    """Returns the terminal (x,y) size for fd.
64
65    Args:
66      fd: The terminal file descriptor.
67
68    Returns:
69      The terminal (x,y) size for fd or None on error.
70    """
71    try:
72      # This magic incantation converts a struct from ioctl(2) containing two
73      # binary shorts to a (rows, columns) int tuple.
74      rc = struct.unpack(b'hh', fcntl.ioctl(fd, termios.TIOCGWINSZ, 'junk'))
75      return (rc[1], rc[0]) if rc else None
76    except:  # pylint: disable=bare-except
77      return None
78
79  xy = _GetXY(0) or _GetXY(1) or _GetXY(2)
80  if not xy:
81    fd = None
82    try:
83      fd = os.open(os.ctermid(), os.O_RDONLY)
84      xy = _GetXY(fd)
85    except:  # pylint: disable=bare-except
86      xy = None
87    finally:
88      if fd is not None:
89        os.close(fd)
90  return xy
91
92
93def _GetTermSizeWindows():
94  """Returns the Windows terminal x and y dimensions."""
95  # pylint:disable=g-import-not-at-top
96  import struct
97  # pylint: disable=g-import-not-at-top
98  from ctypes import create_string_buffer
99  # pylint:disable=g-import-not-at-top
100  from ctypes import windll
101
102  # stdin handle is -10
103  # stdout handle is -11
104  # stderr handle is -12
105
106  h = windll.kernel32.GetStdHandle(-12)
107  csbi = create_string_buffer(22)
108  if not windll.kernel32.GetConsoleScreenBufferInfo(h, csbi):
109    return None
110  (unused_bufx, unused_bufy, unused_curx, unused_cury, unused_wattr,
111   left, top, right, bottom,
112   unused_maxx, unused_maxy) = struct.unpack(b'hhhhHhhhhhh', csbi.raw)
113  x = right - left + 1
114  y = bottom - top + 1
115  return (x, y)
116
117
118def _GetTermSizeEnvironment():
119  """Returns the terminal x and y dimensions from the environment."""
120  return (int(os.environ['COLUMNS']), int(os.environ['LINES']))
121
122
123def _GetTermSizeTput():
124  """Returns the terminal x and y dimemsions from tput(1)."""
125  import subprocess  # pylint: disable=g-import-not-at-top
126  output = encoding.Decode(subprocess.check_output(['tput', 'cols'],
127                                                   stderr=subprocess.STDOUT))
128  cols = int(output)
129  output = encoding.Decode(subprocess.check_output(['tput', 'lines'],
130                                                   stderr=subprocess.STDOUT))
131  rows = int(output)
132  return (cols, rows)
133
134
135_ANSI_CSI = '\x1b'  # ANSI control sequence indicator (ESC)
136_CONTROL_D = '\x04'  # unix EOF (^D)
137_CONTROL_Z = '\x1a'  # Windows EOF (^Z)
138_WINDOWS_CSI_1 = '\x00'  # Windows control sequence indicator #1
139_WINDOWS_CSI_2 = '\xe0'  # Windows control sequence indicator #2
140
141
142def GetRawKeyFunction():
143  """Returns a function that reads one keypress from stdin with no echo.
144
145  Returns:
146    A function that reads one keypress from stdin with no echo or a function
147    that always returns None if stdin does not support it.
148  """
149  # Believe the first helper that doesn't bail.
150  for get_raw_key_function in (_GetRawKeyFunctionPosix,
151                               _GetRawKeyFunctionWindows):
152    try:
153      return get_raw_key_function()
154    except:  # pylint: disable=bare-except
155      pass
156  return lambda: None
157
158
159def _GetRawKeyFunctionPosix():
160  """_GetRawKeyFunction helper using Posix APIs."""
161  # pylint: disable=g-import-not-at-top
162  import tty
163  # pylint: disable=g-import-not-at-top
164  import termios
165
166  def _GetRawKeyPosix():
167    """Reads and returns one keypress from stdin, no echo, using Posix APIs.
168
169    Returns:
170      The key name, None for EOF, <*> for function keys, otherwise a
171      character.
172    """
173    ansi_to_key = {
174        'A': '<UP-ARROW>',
175        'B': '<DOWN-ARROW>',
176        'D': '<LEFT-ARROW>',
177        'C': '<RIGHT-ARROW>',
178        '5': '<PAGE-UP>',
179        '6': '<PAGE-DOWN>',
180        'H': '<HOME>',
181        'F': '<END>',
182        'M': '<DOWN-ARROW>',
183        'S': '<PAGE-UP>',
184        'T': '<PAGE-DOWN>',
185    }
186
187    # Flush pending output. sys.stdin.read() would do this, but it's explicitly
188    # bypassed in _GetKeyChar().
189    sys.stdout.flush()
190
191    fd = sys.stdin.fileno()
192
193    def _GetKeyChar():
194      return encoding.Decode(os.read(fd, 1))
195
196    old_settings = termios.tcgetattr(fd)
197    try:
198      tty.setraw(fd)
199      c = _GetKeyChar()
200      if c == _ANSI_CSI:
201        c = _GetKeyChar()
202        while True:
203          if c == _ANSI_CSI:
204            return c
205          if c.isalpha():
206            break
207          prev_c = c
208          c = _GetKeyChar()
209          if c == '~':
210            c = prev_c
211            break
212        return ansi_to_key.get(c, '')
213    except:  # pylint:disable=bare-except
214      c = None
215    finally:
216      termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
217    return None if c in (_CONTROL_D, _CONTROL_Z) else c
218
219  return _GetRawKeyPosix
220
221
222def _GetRawKeyFunctionWindows():
223  """_GetRawKeyFunction helper using Windows APIs."""
224  # pylint: disable=g-import-not-at-top
225  import msvcrt
226
227  def _GetRawKeyWindows():
228    """Reads and returns one keypress from stdin, no echo, using Windows APIs.
229
230    Returns:
231      The key name, None for EOF, <*> for function keys, otherwise a
232      character.
233    """
234    windows_to_key = {
235        'H': '<UP-ARROW>',
236        'P': '<DOWN-ARROW>',
237        'K': '<LEFT-ARROW>',
238        'M': '<RIGHT-ARROW>',
239        'I': '<PAGE-UP>',
240        'Q': '<PAGE-DOWN>',
241        'G': '<HOME>',
242        'O': '<END>',
243    }
244
245    # Flush pending output. sys.stdin.read() would do this it's explicitly
246    # bypassed in _GetKeyChar().
247    sys.stdout.flush()
248
249    def _GetKeyChar():
250      return encoding.Decode(msvcrt.getch())
251
252    c = _GetKeyChar()
253    # Special function key is a two character sequence; return the second char.
254    if c in (_WINDOWS_CSI_1, _WINDOWS_CSI_2):
255      return windows_to_key.get(_GetKeyChar(), '')
256    return None if c in (_CONTROL_D, _CONTROL_Z) else c
257
258  return _GetRawKeyWindows
259
260
261def ForceEnableAnsi():
262  """Attempts to enable virtual terminal processing on Windows.
263
264  Returns:
265    bool: True if ANSI support is now active; False otherwise.
266  """
267  if platforms.OperatingSystem.Current() != platforms.OperatingSystem.WINDOWS:
268    return False
269
270  try:
271    import ctypes  # pylint:disable=g-import-not-at-top
272
273    enable_virtual_terminal_processing = 0x0004
274    h = ctypes.windll.kernel32.GetStdHandle(-11)  # stdout handle is -11
275    old_mode = ctypes.wintypes.DWORD()
276
277    if ctypes.windll.kernel32.GetConsoleMode(h, ctypes.byref(old_mode)):
278      if ctypes.windll.kernel32.SetConsoleMode(
279          h, old_mode.value | enable_virtual_terminal_processing):
280        return True
281  except (OSError, AttributeError):
282    pass  # If we cannot force ANSI, we should simply return False
283  return False
284