1# -*- coding: utf-8 -*- 

2""" 

3@file 

4@brief Magic command to handle files 

5""" 

6import os 

7from IPython.core.magic import magics_class, line_magic 

8 

9from .magic_class import MagicClassWithHelpers 

10from .magic_parser import MagicCommandParser 

11from ..filehelper import encrypt_stream, decrypt_stream 

12 

13 

14@magics_class 

15class MagicCrypt(MagicClassWithHelpers): 

16 

17 """ 

18 Defines magic commands to encrypt and decrypt file. 

19 """ 

20 

21 @staticmethod 

22 def endecrypt_file_parser(encrypt): 

23 """ 

24 defines the way to parse the magic command ``%encrypt_file`` and ``%decrypt_file`` 

25 

26 @param encrypt True to encrypt or False to decrypt 

27 @return parser 

28 """ 

29 task = "encrypt" if encrypt else "decrypt" 

30 parser = MagicCommandParser(prog="%scrypt_file" % task[:2], 

31 description='%s a file' % task + 

32 '\ndoes not work well in Python 2.7 with pycryptodomex') 

33 parser.add_argument( 

34 'source', 

35 help='file to %s' % task) 

36 parser.add_argument( 

37 'dest', 

38 help='location of the %sed file' % task) 

39 parser.add_argument( 

40 'password', 

41 help='password, usually an ascii string with 16x characters') 

42 

43 return parser 

44 

45 @staticmethod 

46 def encrypt_file_parser(): 

47 """ 

48 defines the way to parse the magic command ``%encrypt_file`` 

49 """ 

50 return MagicCrypt.endecrypt_file_parser(True) 

51 

52 @staticmethod 

53 def decrypt_file_parser(): 

54 """ 

55 defines the way to parse the magic command ``%decrypt_file`` 

56 """ 

57 return MagicCrypt.endecrypt_file_parser(False) 

58 

59 @line_magic 

60 def encrypt_file(self, line): 

61 """ 

62 .. nbref:: 

63 :title: %encrypt_file 

64 :tag: nb 

65 :lid: l-nb-encrypt-file 

66 

67 The magic command is equivalent to:: 

68 

69 from pyquickhelper.filehelper import encrypt_stream 

70 

71 password = "password" 

72 source = "file source" 

73 dest = "file destination" 

74 

75 if isinstance(password, str): 

76 password = bytes(password, encoding="ascii") 

77 

78 encrypt_stream(key=password, filename=source, out_filename=dest, 

79 chunksize=os.stat(source).st_size * 2 + 1) 

80 """ 

81 parser = self.get_parser( 

82 MagicCrypt.encrypt_file_parser, "encrypt_file") 

83 args = self.get_args(line, parser) 

84 

85 if args is not None: 

86 password = args.password 

87 source = args.source 

88 dest = args.dest 

89 

90 if isinstance(password, str): 

91 password = bytes(password, encoding="ascii") 

92 

93 return encrypt_stream(key=password, filename=source, out_filename=dest, 

94 chunksize=os.stat(source).st_size * 2 + 1) 

95 return None # pragma: no cover 

96 

97 @line_magic 

98 def decrypt_file(self, line): 

99 """ 

100 .. nbref:: 

101 :title: %decrypt_file 

102 

103 The magic command is equivalent to:: 

104 

105 from pyquickhelper.filehelper import decrypt_stream 

106 

107 password = "password" 

108 source = "file source" 

109 dest = "file destination" 

110 

111 if isinstance(password, str): 

112 password = bytes(password, encoding="ascii") 

113 

114 decrypt_stream(key=password, filename=source, out_filename=dest, 

115 chunksize=os.stat(source).st_size * 2 + 1) 

116 """ 

117 parser = self.get_parser( 

118 MagicCrypt.decrypt_file_parser, "decrypt_file") 

119 args = self.get_args(line, parser) 

120 

121 if args is not None: 

122 password = args.password 

123 source = args.source 

124 dest = args.dest 

125 

126 if isinstance(password, str): 

127 password = bytes(password, encoding="ascii") 

128 

129 return decrypt_stream(key=password, filename=source, out_filename=dest, 

130 chunksize=os.stat(source).st_size * 2 + 1) 

131 return None # pragma: no cover 

132 

133 

134def register_file_magics(ip=None): # pragma: no cover 

135 """ 

136 Register magics function, can be called from a notebook. 

137 

138 @param ip from ``get_ipython()`` 

139 """ 

140 if ip is None: 

141 from IPython import get_ipython 

142 ip = get_ipython() 

143 ip.register_magics(MagicCrypt)