diff --git a/config/aaa.py b/config/aaa.py index d8f38644..3c761871 100644 --- a/config/aaa.py +++ b/config/aaa.py @@ -149,68 +149,45 @@ def login(auth_protocol): add_table_kv('AAA', 'authentication', 'login', val) authentication.add_command(login) - # cmd: aaa authorization @click.command() -@click.argument('protocol', nargs=-1, type=click.Choice([ "tacacs+", "local", "default"])) +@click.argument('protocol', nargs=-1, type=click.Choice([ "tacacs+", "local", "tacacs+ local"])) def authorization(protocol): - """ - Switch AAA authorization default - authorization {tacacs+ | local} [local] - """ + """Switch AAA authorization [tacacs+ | local | '\"tacacs+ local\"']""" if len(protocol) == 0: click.echo('Argument "protocol" is required') return - if len(protocol) > 2: - click.fail('Only two "protocol" can be set at most') - - if 'default' in protocol: - if len(protocol) >= 2: - click.fail('Invalidated parameter') - - del_table_key('AAA', 'authorization', 'login') - return - - if protocol[0] == 'local' and len(protocol) != 1: - click.fail('Invalidated parameter') - - if len(protocol) != len(set(protocol)): - click.fail('Invalidated parameter') - - add_table_kv('AAA', 'authorization', 'login', ','.join(protocol)) + if len(protocol) == 1 and (protocol[0] == 'tacacs+' or protocol[0] == 'local'): + add_table_kv('AAA', 'authorization', 'login', protocol[0]) + elif len(protocol) == 1 and protocol[0] == 'tacacs+ local': + add_table_kv('AAA', 'authorization', 'login', 'tacacs+,local') + else: + click.echo('Not a valid command') aaa.add_command(authorization) - # cmd: aaa accounting @click.command() -@click.argument('protocol', nargs=-1, type=click.Choice(["disable", "tacacs+", "local"])) +@click.argument('protocol', nargs=-1, type=click.Choice(["disable", "tacacs+", "local", "tacacs+ local"])) def accounting(protocol): - """ - Switch AAA accounting disable - accounting {tacacs+ | local} [tacacs+ | local] - """ + """Switch AAA accounting [disable | tacacs+ | local | '\"tacacs+ local\"']""" if len(protocol) == 0: click.echo('Argument "protocol" is required') return - if len(protocol) > 2: - click.fail('Only two "protocol" can be set at most') - - if 'disable' in protocol: - if len(protocol) >= 2: - click.fail('Invalidated parameter') - - del_table_key('AAA', 'accounting', 'login') - return - - if len(protocol) != len(set(protocol)): - click.fail('Invalidated parameter') - - add_table_kv('AAA', 'accounting', 'login', ','.join(protocol)) + if len(protocol) == 1: + if protocol[0] == 'tacacs+' or protocol[0] == 'local': + add_table_kv('AAA', 'accounting', 'login', protocol[0]) + elif protocol[0] == 'tacacs+ local': + add_table_kv('AAA', 'accounting', 'login', 'tacacs+,local') + elif protocol[0] == 'disable': + del_table_key('AAA', 'accounting', 'login') + else: + click.echo('Not a valid command') + else: + click.echo('Not a valid command') aaa.add_command(accounting) - @click.group() def tacacs(): """TACACS+ server configuration""" @@ -428,9 +405,8 @@ def sourceip(ctx, src_ip): click.echo('Invalid ip address') return - v6_invalid_list = [ipaddress.IPv6Address(str('0::0')), ipaddress.IPv6Address(str('0::1'))] - net = ipaddress.ip_network(str(src_ip), strict=False) - + v6_invalid_list = [ipaddress.IPv6Address('0::0'), ipaddress.IPv6Address('0::1')] + net = ipaddress.ip_network(src_ip, strict=False) if (net.version == 4): if src_ip == "0.0.0.0": click.echo('enter non-zero ip address') @@ -470,9 +446,8 @@ def nasip(ctx, nas_ip): click.echo('Invalid ip address') return - v6_invalid_list = [ipaddress.IPv6Address(str('0::0')), ipaddress.IPv6Address(str('0::1'))] - net = ipaddress.ip_network(str(nas_ip), strict=False) - + v6_invalid_list = [ipaddress.IPv6Address('0::0'), ipaddress.IPv6Address('0::1')] + net = ipaddress.ip_network(nas_ip, strict=False) if (net.version == 4): if nas_ip == "0.0.0.0": click.echo('enter non-zero ip address') diff --git a/tests/aaa_test.py b/tests/aaa_test.py index ba22a3c6..e554937e 100644 --- a/tests/aaa_test.py +++ b/tests/aaa_test.py @@ -230,7 +230,7 @@ def test_config_aaa_tacacs(self, get_cmd_module): # test tacacs + local authorization result = runner.invoke(config.config.commands["aaa"],\ - ["authorization", "tacacs+", "local"], obj=db) + ["authorization", "tacacs+ local"], obj=db) print(result.exit_code) print(result.output) assert result.exit_code == 0 @@ -258,7 +258,7 @@ def test_config_aaa_tacacs(self, get_cmd_module): # test tacacs + local accounting result = runner.invoke(config.config.commands["aaa"],\ - ["accounting", "tacacs+", "local"], obj=db) + ["accounting", "tacacs+ local"], obj=db) print(result.exit_code) print(result.output) assert result.exit_code == 0