Skip to content

Commit

Permalink
feat(secure_api_key): implement orgazization GUI to secure API key
Browse files Browse the repository at this point in the history
  • Loading branch information
@clementb49
clementb49 committed Aug 24, 2024
1 parent 9284661 commit c1d97c6
Showing 1 changed file with 123 additions and 102 deletions.
225 changes: 123 additions & 102 deletions basilisk/gui/account_dialog.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,19 +8,19 @@
Account,
AccountOrganization,
AccountSource,
ApiKeyStorageMethodEnum,
KeyStorageMethodEnum,
get_account_source_labels,
)
from basilisk.config import conf
from basilisk.provider import get_provider, providers

log = getLogger(__name__)

api_storage_methods = {
key_storage_methods = {
# Translators: A label for the API key storage method in the account dialog
ApiKeyStorageMethodEnum.plain: _("Plain text"),
KeyStorageMethodEnum.plain: _("Plain text"),
# Translators: A label for the API key storage method in the account dialog
ApiKeyStorageMethodEnum.system: _("System keyring"),
KeyStorageMethodEnum.system: _("System keyring"),
}


Expand All @@ -47,12 +47,35 @@ def init_ui(self):
sizer = wx.BoxSizer(wx.VERTICAL)
panel.SetSizer(sizer)

label = wx.StaticText(panel, label=_("&Name:"), style=wx.ALIGN_LEFT)
label = wx.StaticText(
panel,
# Translators: A label in account dialog
label=_("&Name:"),
style=wx.ALIGN_LEFT,
)
sizer.Add(label, 0, wx.ALL, 5)
self.name = wx.TextCtrl(panel)
sizer.Add(self.name, 0, wx.EXPAND)

label = wx.StaticText(panel, label=_("API &Key:"), style=wx.ALIGN_LEFT)
label = wx.StaticText(
panel,
# Translators: A label in account dialog
label=_("Organisation key storage &method:"),
style=wx.ALIGN_LEFT,
)
sizer.Add(label, 0, wx.ALL, 5)
self.key_storage_method = wx.ComboBox(
panel,
choices=list(key_storage_methods.values()),
style=wx.CB_READONLY,
)
sizer.Add(self.key_storage_method, 0, wx.EXPAND)
label = wx.StaticText(
panel,
# Translators: A label in account dialog
label=_("Organisation &Key:"),
style=wx.ALIGN_LEFT,
)
sizer.Add(label, 0, wx.ALL, 5)
self.key = wx.TextCtrl(panel)
sizer.Add(self.key, 0, wx.EXPAND)
Expand All @@ -71,9 +94,18 @@ def init_ui(self):
sizer.Add(bSizer, 0, wx.ALL, 5)

def init_data(self):
if self.organization:
self.name.SetValue(self.organization.name)
self.key.SetValue(self.organization.key.get_secret_value())
if not self.organization:
self.key_storage_method.SetSelection(0)
return

self.name.SetValue(self.organization.name)
index = -1
for i, method in enumerate(key_storage_methods.keys()):
if method == self.organization.key_storage_method:
index = i
break
self.key_storage_method.SetSelection(index)
self.key.SetValue(self.organization.key.get_secret_value())

def update_data(self):
pass
Expand All @@ -84,17 +116,28 @@ def on_ok(self, event):
wx.MessageBox(msg, _("Error"), wx.OK | wx.ICON_ERROR)
self.name.SetFocus()
return
if self.key_storage_method.GetSelection() == -1:
msg = _("Please select a key storage method")
wx.MessageBox(msg, _("Error"), wx.OK | wx.ICON_ERROR)
self.key_storage_method.SetFocus()
return
if not self.key.GetValue():
msg = _("Please enter a key")
wx.MessageBox(msg, _("Error"), wx.OK | wx.ICON_ERROR)
self.key.SetFocus()
return
key_storage_method = list(key_storage_methods.keys())[
self.key_storage_method.GetSelection()
]
if self.organization:
self.organization.name = self.name.GetValue()
self.organization.key_storage_method = (key_storage_method,)
self.organization.key = SecretStr(self.key.GetValue())
else:
self.organization = AccountOrganization(
name=self.name.GetValue(), key=self.key.GetValue()
name=self.name.GetValue(),
key_storage_method=key_storage_method,
key=SecretStr(self.key.GetValue()),
)
self.EndModal(wx.ID_OK)

Expand Down Expand Up @@ -144,6 +187,7 @@ def init_ui(self):
self.organization_list.Bind(
wx.EVT_LIST_ITEM_SELECTED, self.on_item_selected
)
self.organization_list.Bind(wx.EVT_KEY_DOWN, self.on_org_list_key_down)
sizer.Add(self.organization_list, 1, wx.EXPAND)

add_btn = wx.Button(panel, label=_("&Add"))
Expand Down Expand Up @@ -187,10 +231,6 @@ def update_data(self):
)

def update_ui(self):
if not self.organizations:
self.edit_btn.Disable()
self.remove_btn.Disable()
return
selected_item = self.organization_list.GetFirstSelected()
if selected_item == -1:
self.edit_btn.Disable()
Expand All @@ -205,18 +245,7 @@ def update_ui(self):
self.remove_btn.Enable()

def on_item_selected(self, event):
selected_item = self.organization_list.GetFirstSelected()
if selected_item == -1:
self.edit_btn.Disable()
self.remove_btn.Disable()
return
organization = self.organizations[selected_item]
if organization.source == AccountSource.ENV_VAR:
self.edit_btn.Disable()
self.remove_btn.Disable()
return
self.edit_btn.Enable()
self.remove_btn.Enable()
self.update_ui()

def on_add(self, event):
dialog = EditAccountOrganizationDialog(self, _("Add organization"))
Expand Down Expand Up @@ -251,13 +280,11 @@ def on_edit(self, event):
if dialog.ShowModal() == wx.ID_OK:
organization = dialog.organization
self.organizations[selected_item] = organization
self.organization_list.SetStringItem(
selected_item, 0, organization.name
)
self.organization_list.SetStringItem(
self.organization_list.SetItem(selected_item, 0, organization.name)
self.organization_list.SetItem(
selected_item, 1, organization.key.get_secret_value()
)
self.organization_list.SetStringItem(
self.organization_list.SetItem(
selected_item,
2,
self.account_source_labels.get(
Expand All @@ -273,19 +300,20 @@ def on_edit(self, event):
self.organization_list.EnsureVisible(selected_item)

def on_remove(self, event):
item = self.organization_list.GetFirstSelected()
organization_id = self.organizations[item].id
organization_name = self.organizations[item].name
index = self.organization_list.GetFirstSelected()
organization = self.organizations[index]
# Translators: A confirmation message in account dialog for removing organization
msg = _("Are you sure you want to remove the organization {}?").format(
organization_name
organization.name
)
if wx.MessageBox(msg, _("Confirmation"), wx.YES_NO) != wx.YES:
return
self.organizations.pop(item)
self.organization_list.DeleteItem(item)
if self.account.active_organization_id == organization_id:
organization.delete_keyring_password()
self.organization_list.Select(index - 1)
self.organization_list.DeleteItem(index)
if self.account.active_organization_id == organization.id:
self.account.active_organization_id = None
self.organizations.pop(index)
self.update_ui()

def onOK(self, event):
Expand All @@ -295,6 +323,14 @@ def onOK(self, event):
def onCancel(self, event):
self.EndModal(wx.ID_CANCEL)

def on_org_list_key_down(self, event: wx.KeyEvent):
if event.GetKeyCode() == wx.WXK_RETURN:
self.on_edit(event)
elif event.GetKeyCode() == wx.WXK_DELETE:
self.on_remove(event)
else:
event.Skip()


class EditAccountDialog(wx.Dialog):
def __init__(self, parent, title, size=(400, 400), account: Account = None):
Expand Down Expand Up @@ -347,7 +383,7 @@ def init_ui(self):
sizer.Add(label, 0, wx.ALL, 5)
self.api_key_storage_method = wx.ComboBox(
panel,
choices=list(api_storage_methods.values()),
choices=list(key_storage_methods.values()),
style=wx.CB_READONLY,
)
sizer.Add(self.api_key_storage_method, 0, wx.EXPAND)
Expand Down Expand Up @@ -385,55 +421,41 @@ def init_ui(self):
sizer.Add(bSizer, 0, wx.ALL, 5)

def init_data(self):
if self.account:
self.name.SetValue(self.account.name)
if not self.account:
self.api_key_storage_method.SetSelection(0)
return
self.name.SetValue(self.account.name)
index = -1
for i, provider in enumerate(providers):
if provider.name == self.account.provider.name:
index = i
break
self.provider.SetSelection(index)
if self.account.api_key and self.account.api_key_storage_method:
index = -1
for i, provider in enumerate(providers):
if provider.name == self.account.provider.name:
for i, method in enumerate(key_storage_methods.keys()):
if method == self.account.api_key_storage_method:
index = i
break
self.provider.SetSelection(index)
if self.account.api_key and self.account.api_key_storage_method:
index = -1
for i, method in enumerate(api_storage_methods.keys()):
if method == self.account.api_key_storage_method:
index = i
break
self.api_key_storage_method.SetSelection(index)
self.api_key.SetValue(self.account.api_key.get_secret_value())
self.organization.Enable(
self.account.provider.organization_mode_available
)
if not self.account.provider.organization_mode_available:
return
if self.account.organizations:
choices = [_("Personal")] + [
organization.name
for organization in self.account.organizations
]
self.organization.SetItems(choices)
if self.account.active_organization_id:
index = -1
for i, organization in enumerate(self.account.organizations):
if organization.id == self.account.active_organization_id:
index = i + 1
break
self.organization.SetSelection(index)
else:
self.organization.SetSelection(0)

if self.account.active_organization_id:
index = -1
for i, organization in enumerate(
self.account.organizations
):
if (
organization.id
== self.account.active_organization_id
):
index = i + 1
break
self.organization.SetSelection(index)
self.api_key_storage_method.SetSelection(index)
self.api_key.SetValue(self.account.api_key.get_secret_value())
self.organization.Enable(
self.account.provider.organization_mode_available
)
if not self.account.provider.organization_mode_available:
return
if self.account.organizations:
choices = [_("Personal")] + [
organization.name for organization in self.account.organizations
]
self.organization.SetItems(choices)
if self.account.active_organization_id:
index = -1
for i, organization in enumerate(self.account.organizations):
if organization.id == self.account.active_organization_id:
index = i + 1
break
self.organization.SetSelection(index)

def update_ui(self):
provider_index = self.provider.GetSelection()
Expand Down Expand Up @@ -481,28 +503,27 @@ def on_ok(self, event):
active_organization = self.account.organizations[
organization_index - 1
].id
api_key_storage_method = None
api_key = None
if provider.require_api_key:
api_key_storage_method = list(key_storage_methods.keys())[
self.api_key_storage_method.GetSelection()
]
api_key = SecretStr(self.api_key.GetValue())
if self.account:
self.account.name = self.name.GetValue()
self.account.provider = provider
if provider.require_api_key:
self.account.api_key_storage_method = list(
api_storage_methods.keys()
)[self.api_key_storage_method.GetSelection()]
self.account.api_key = SecretStr(self.api_key.GetValue())
self.account.api_key_storage_method = api_key_storage_method
self.account.api_key = api_key
self.account.active_organization_id = active_organization
else:
self.account = Account.model_validate(
dict(
name=self.name.GetValue(),
provider=provider,
api_key_storage_method=list(api_storage_methods.keys())[
self.api_key_storage_method.GetSelection()
],
api_key=SecretStr(self.api_key.GetValue()),
active_organization_id=active_organization,
source=AccountSource.CONFIG,
),
context={"new_account": True},
self.account = Account(
name=self.name.GetValue(),
provider=provider,
api_key_storage_method=api_key_storage_method,
api_key=api_key,
active_organization_id=active_organization,
source=AccountSource.CONFIG,
)
self.EndModal(wx.ID_OK)

Expand Down

0 comments on commit c1d97c6

Please sign in to comment.