ServalShell/shell.py at master · em-1001/ServalShell · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import os
import subprocess
from translate import translate
import sys
sys.path.append("./Tellina")
from bashlint.data_tools import bash_tokenizer, bash_parser, ast2tokens, ast2command
from nlp_tools import tokenizer
from bashlint import data_tools
from encoder_decoder import slot_filling

import warnings
warnings.filterwarnings("ignore")


def post_processing(nl, _bash):
    node = bash_parser(_bash)
    _, _, nl_filler= tokenizer.ner_tokenizer(nl)[1]
    slot_filling.heuristic_slot_filling(node, nl_filler)

    bash = ast2command(node)

    bash2 = slot_filling.stupid_slot_matching(nl, _bash)

    if _bash == bash:
        bash = bash2

    return bash


def servalshell():
    serval_cat = """
       _
       \\`*-.
        )  _`-.
       .  : `. .
       : _   '  \\
       ; *` _.   `*-_.
       `-.-'          `-.
         ;       `       `.
         :.       .        \\
         . \\  .   :   .-'   .
         '  `+.;  ;  '      :
         :  '  |    ;       ;-.
         ; '   : :`-:     _.`* ;
      .*' /  .*' ; .*`- +'  `*'
      `*-*   `*-*  `*-*'
    """
    print('\033[92m' + serval_cat)
    prompt = '\033[32m' + 'ServalShell' + '\033[92m' + ':~$ '
    while True:
        recommend_flag = 0

        nl = input(prompt)

        nl_split = nl.split(' ')

        # Execute bash command directly
        if nl_split[0] == '-d' or nl_split[0] == '--direct':
          direct_bash = ' '.join(nl_split[1:])
          try:
            output = subprocess.check_output(direct_bash, shell=True, text=True)
            print(output)
          except subprocess.CalledProcessError as e:
            print(e)
            print("")
          continue

        # Describes usage and options
        elif nl_split[0] == '-h' or nl_split[0] == '--help':
          print("If you enter a command in natural language, the program automatically translates it into a bash command and executes it.")
          print("If execution fails because the bash command translated by the model is incorrect, It will recommend several command structures.")
          print("Additionally, the following options are available.\n")
          print("-d [cmd],  --direct [cmd]              Execute bash command directly")
          print("-r [nl],  --recommend [nl]             Even if the command execution is successful, Recommended Command Structure is displayed")
          print("-h,  --help                            Describes usage and options")
          print("-q,  --quit                            Quit Servalshell")
          print("")
          continue

        # Even if the command execution is successful, Recommended Command Structure is displayed
        elif nl_split[0] == '-r' or nl_split[0] == '--recommend':
            recommend_flag = 1
            nl = ' '.join(nl_split[1:])

        # Quit Servalshell
        elif nl_split[0] == '-q' or nl_split[0] == '--quit':
          break

        nl_preprocess = ' '.join(tokenizer.ner_tokenizer(nl)[0])
        _bash, _ = translate(nl_preprocess)

        bash_list = []
        for cmd in _bash:
            try:
                bash = post_processing(nl, cmd)
                bash_list.append(bash)
            except:
                continue

        if len(bash_list) == 0:
            print("Failed at Abstract syntax tree...")
            print("")
            continue

        print("translated bash: " + str(bash_list[0]))
        try:
            output = subprocess.check_output(bash_list[0], shell=True, text=True)
            print(output)

            if recommend_flag == 1:
                print("Recommended Command Structure" + "\033[92m")
                for rcs in _bash:
                    print(rcs)
                print("")

        except subprocess.CalledProcessError as e:
            print("\nRecommended Command Structure" + "\033[92m")
            for rcs in _bash:
                print(rcs)
            print("")


if __name__ == '__main__':
    servalshell()