1"""
2Functions for working with the codepage on Windows systems
3"""
4
5import logging
6from contextlib import contextmanager
7
8from salt.exceptions import CodePageError
9
10log = logging.getLogger(__name__)
11
12try:
13    import pywintypes
14    import win32console
15
16    HAS_WIN32 = True
17except ImportError:
18    HAS_WIN32 = False
19
20
21# Although utils are often directly imported, it is also possible to use the loader.
22def __virtual__():
23    """
24    Only load if Win32 Libraries are installed
25    """
26    if not HAS_WIN32:
27        return False, "This utility requires pywin32"
28
29    return "win_chcp"
30
31
32@contextmanager
33def chcp(page_id, raise_error=False):
34    """
35    Gets or sets the codepage of the shell.
36
37    Args:
38
39        page_id (str, int):
40            A number representing the codepage.
41
42        raise_error (bool):
43            ``True`` will raise an error if the codepage fails to change.
44            ``False`` will suppress the error
45
46    Returns:
47        int: A number representing the codepage
48
49    Raises:
50        CodePageError: On unsuccessful codepage change
51    """
52    if not isinstance(page_id, int):
53        try:
54            page_id = int(page_id)
55        except ValueError:
56            error = "The `page_id` needs to be an integer, not {}".format(type(page_id))
57            if raise_error:
58                raise CodePageError(error)
59            log.error(error)
60            return -1
61
62    previous_page_id = get_codepage_id(raise_error=raise_error)
63
64    if page_id and previous_page_id and page_id != previous_page_id:
65        set_code_page = True
66    else:
67        set_code_page = False
68
69    try:
70        if set_code_page:
71            set_codepage_id(page_id, raise_error=raise_error)
72
73        # Subprocesses started from now will use the set code page id
74        yield
75    finally:
76        if set_code_page:
77            # Reset to the old code page
78            set_codepage_id(previous_page_id, raise_error=raise_error)
79
80
81def get_codepage_id(raise_error=False):
82    """
83    Get the currently set code page on windows
84
85    Args:
86
87        raise_error (bool):
88            ``True`` will raise an error if the codepage fails to change.
89            ``False`` will suppress the error
90
91    Returns:
92        int: A number representing the codepage
93
94    Raises:
95        CodePageError: On unsuccessful codepage change
96    """
97    try:
98        return win32console.GetConsoleCP()
99    except pywintypes.error as exc:
100        _, _, msg = exc.args
101        error = "Failed to get the windows code page: {}".format(msg)
102        if raise_error:
103            raise CodePageError(error)
104        else:
105            log.error(error)
106        return -1
107
108
109def set_codepage_id(page_id, raise_error=False):
110    """
111    Set the code page on windows
112
113    Args:
114
115        page_id (str, int):
116            A number representing the codepage.
117
118        raise_error (bool):
119            ``True`` will raise an error if the codepage fails to change.
120            ``False`` will suppress the error
121
122    Returns:
123        int: A number representing the codepage
124
125    Raises:
126        CodePageError: On unsuccessful codepage change
127    """
128    if not isinstance(page_id, int):
129        try:
130            page_id = int(page_id)
131        except ValueError:
132            error = "The `page_id` needs to be an integer, not {}".format(type(page_id))
133            if raise_error:
134                raise CodePageError(error)
135            log.error(error)
136            return -1
137    try:
138        win32console.SetConsoleCP(page_id)
139        return get_codepage_id(raise_error=raise_error)
140    except pywintypes.error as exc:
141        _, _, msg = exc.args
142        error = "Failed to set the windows code page: {}".format(msg)
143        if raise_error:
144            raise CodePageError(error)
145        else:
146            log.error(error)
147        return -1
148