Coverage for kryptic_cypher/app.py: 84%

1""" 

2Module that contains the click entrypoint for our cli interface. 

3 

4This is currently only for encoding and decoding data using the encode and decode commands.  

5""" 

6 

7import base64 

8from enum import Enum 

9import functools 

10from io import BytesIO 

11import shutil 

12from typing import IO, Callable 

13import click 

14from .cypher import Cypher, CypherWithKey, registered_cyphers 

15 

16 

17@click.group() 

18@click.pass_context 

19def main(ctx: click.Context): 

20 """ 

21 main group that represents the top-level: ***zombie-nomnom*** 

22 

23 This will be used to decorate sub-commands for zombie-nomnom. 

24 

25 ***Example Usage:*** 

26 ```python 

27 @main.command("sub-command") 

28 def sub_command(): 

29 # do actual meaningful work. 

30 pass 

31 ``` 

32 """ 

33 

34 

35def resolve_cypher( 

36 cypher: str, 

37 text: str, 

38 input: str, 

39 key: str, 

40) -> Cypher | CypherWithKey: 

41 if not text and not input: 

42 raise click.ClickException("You must specify either -t or -i") 

43 

44 cypher_instance = registered_cyphers[cypher] 

45 if isinstance(cypher_instance, CypherWithKey): 

46 if not key: 

47 raise click.ClickException("You must specify -k") 

48 result = cypher_instance.validate_key(key) 

49 if not result.success: 

50 raise click.ClickException("\n".join(result.messages)) 

51 

52 return cypher_instance 

53 

54 

55def process_output(output: str | None, encoded_text: IO): 

56 if output: 

57 flags = "wb" if getattr(encoded_text, "mode", None) == "b" else "w" 

58 with open(output, flags) as f: 

59 shutil.copyfileobj(encoded_text, f) 

60 else: 

61 if getattr(encoded_text, "mode", None) == "b": 

62 full_value = encoded_text.read() 

63 binary_string = base64.b64encode(full_value).decode("utf-8") 

64 click.echo(binary_string) 

65 else: 

66 click.echo(encoded_text.read()) 

67 

68 

69@main.command("encode") 

70@click.option( 

71 "-c", 

72 "--cypher", 

73 help="The cypher to use", 

74 required=True, 

75 type=click.Choice(list(registered_cyphers.keys())), 

76) 

77@click.option("-t", "--text", help="The text to encode", required=False) 

78@click.option("-k", "--key", help="The input file to read text from", required=False) 

79@click.option("-i", "--input", help="The input file to read text from", required=False) 

80@click.option( 

81 "-o", 

82 "--output", 

83 help="The output file to write text to", 

84 required=False, 

85 type=click.Path(writable=True, dir_okay=False), 

86) 

87@click.option( 

88 "-b", 

89 "--binary", 

90 help="The output file to write text to", 

91 required=False, 

92 is_flag=True, 

93) 

94def encode( 

95 cypher: str, 

96 text: str, 

97 input: str, 

98 binary: bool, 

99 key: str, 

100 output: str | None, 

101): 

102 """ 

103 CLI command to encode text using a cypher in our system that will check to make sure the usage is valid i.e. input is given and key is valid if key is required. 

104 """ 

105 cypher_instance = resolve_cypher(cypher, text, input, key) 

106 

107 if input: 

108 with open(input, "rb" if binary else "r") as f: 

109 text = f.read() 

110 

111 if isinstance(cypher_instance, CypherWithKey): 

112 encoded_text = cypher_instance.encode(text, key) 

113 else: 

114 encoded_text = cypher_instance.encode(text) 

115 

116 process_output(output, encoded_text) 

117 

118 

119@main.command("decode") 

120@click.option( 

121 "-c", 

122 "--cypher", 

123 help="The cypher to use", 

124 required=True, 

125 type=click.Choice(list(registered_cyphers.keys())), 

126) 

127@click.option("-t", "--text", help="The text to encode", required=False) 

128@click.option("-i", "--input", help="The input file to read text from", required=False) 

129@click.option("-k", "--key", help="The input file to read text from", required=False) 

130@click.option( 

131 "-o", 

132 "--output", 

133 help="The output file to write text to", 

134 required=False, 

135 type=click.Path(writable=True, dir_okay=False), 

136) 

137@click.option( 

138 "-b", 

139 "--binary", 

140 help="The output file to write text to", 

141 required=False, 

142 is_flag=True, 

143) 

144def decode( 

145 cypher: str, 

146 text: str, 

147 input: str, 

148 key: str, 

149 output: str, 

150 binary: bool, 

151): 

152 """ 

153 CLI command to decode text using a cypher in our system that will check to make sure the usage is valid i.e. input is given and key is valid if key is required. 

154 """ 

155 cypher_instance = resolve_cypher(cypher, text, input, key) 

156 if input: 

157 with open(input, "rb" if binary else "r") as f: 

158 text = f.read() 

159 

160 if isinstance(cypher_instance, CypherWithKey): 

161 encoded_text = cypher_instance.decode(text, key) 

162 else: 

163 encoded_text = cypher_instance.decode(text) 

164 

165 process_output(output, encoded_text)