From 292360ae4f7f0e6ce4df651cbf70f5dedf553bbf Mon Sep 17 00:00:00 2001
From: gurland <stas.gurland@gmail.com>
Date: Sat, 4 Jul 2020 23:16:25 +0300
Subject: [PATCH] Fix #13: validate group index and change group type selection

---
 cleaner.py | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/cleaner.py b/cleaner.py
index da587ee..e781550 100644
--- a/cleaner.py
+++ b/cleaner.py
@@ -42,21 +42,22 @@ def get_all_dialogs():
         return dialogs
 
     def select_supergroup(self):
-        dialogs = self.get_all_dialogs()
+        group_types = {
+            '1': 'supergroup',
+            '2': 'group'
+        }
 
-        print('1. Supergroup\n2. (non super)Group')
-        group_type_n = int(input('Insert group type: '))
-        print('')
+        dialogs = self.get_all_dialogs()
 
-        if group_type_n == 1:
-            self.group_type = 'supergroup'
-        elif group_type_n == 2:
-            self.group_type = 'group'
-        else:
+        print('\n'.join((f'{i}. {name.capitalize()}' for i, name in group_types.items())))
+        try:
+            self.group_type = group_types[input('Insert group type number: ')]
+        except KeyError:
             print('Invalid group type. Exiting..')
-            exit()
+            exit(-1)
+        print('')
 
-        groups = [x for x in dialogs if x.chat.type == self.group_type ]
+        groups = [x for x in dialogs if x.chat.type == self.group_type]
 
         for i, group in enumerate(groups):
             print(f'{i+1}. {group.chat.title}')
@@ -64,6 +65,10 @@ def select_supergroup(self):
         print('')
 
         group_n = int(input('Insert group number: '))
+        if group_n not in range(1, len(groups)+1):
+            print('Invalid group number. Exiting...')
+            exit(-1)
+
         selected_group = groups[group_n - 1]
 
         selected_group_peer = app.resolve_peer(selected_group.chat.id)