Skip to content

Commit

Permalink
fix: data lose after pt update
Browse files Browse the repository at this point in the history
  • Loading branch information
@KawaiiZapic
KawaiiZapic committed Feb 3, 2024
1 parent 934ccc6 commit af361df
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 98 deletions.
73 changes: 17 additions & 56 deletions Config.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,74 +7,36 @@
namespace Community.PowerToys.Run.Plugin.TOTP {
public class OTPList {
public class KeyEntry {
public string Name { get; set; }
public string Key { get; set; }
public bool IsEncrypted { get; set; }
public string Name = "";
public string Key = "";
public bool IsEncrypted = false;

}

public int Version { get; set; }
public List<KeyEntry> Entries { get; set; }
public int Version = 2;
public List<KeyEntry> Entries = new();
}


public static class Config {

public static List<OTPList.KeyEntry> LoadKeyList() {
var nc = new PluginJsonStorage<OTPList>();
var config = nc.Load();
if (config.Entries == null) {
config.Version = 2;
config.Entries = new List<OTPList.KeyEntry>();
nc.Save();
}
return config.Entries;
}

public static void SaveKeyList(List<OTPList.KeyEntry> list) {
var nc = new PluginJsonStorage<OTPList>();
var config = nc.Load();
foreach (var entry in list) {
if (entry.IsEncrypted != true) {
entry.Key = EncryptKey(entry.Key);
entry.IsEncrypted = true;
}
}
config.Entries = list;
nc.Save();
}

public static string DecryptKey(string encrypted) {
return Encoding.UTF8.GetString(ProtectedData.Unprotect(Convert.FromBase64String(encrypted), null, DataProtectionScope.CurrentUser));
}

public static string EncryptKey(string unencrypted) {
return Convert.ToBase64String(ProtectedData.Protect(Encoding.UTF8.GetBytes(unencrypted), null, DataProtectionScope.CurrentUser));
}
}


public static class ConfigMigratorV0 {
private static readonly string DataDirectoryV0 = Environment.ExpandEnvironmentVariables("%LOCALAPPDATA%") + "\\Microsoft\\PowerToys\\PowerToys Run\\Settings\\Plugins\\Zapic.Plugin.TOTP\\";
private static readonly string ConfigPathV0 = DataDirectoryV0 + "TOTPList.json";
private static readonly string DataDirectoryV1 = Environment.ExpandEnvironmentVariables("%LOCALAPPDATA%") + "\\Microsoft\\PowerToys\\PowerToys Run\\Settings\\Plugins\\TOTP\\";
private static readonly string ConfigPathV1 = DataDirectoryV1 + "Config.json";

public class StructV0 {
public string Name { get; set; }
public string Key { get; set; }

public string Name = "";
public string Key = "";
}

public class StructV1 {
public class KeyEntry {
public string Name { get; set; }
public string Key { get; set; }
public bool IsEncrypted { get; set; }
public string Name = "";
public string Key = "";
public bool IsEncrypted = false;

}
public int Version { get; set; }
public List<KeyEntry> Entries { get; set; }
public int Version = 1;
public List<KeyEntry> Entries = new();

}

Expand Down Expand Up @@ -109,14 +71,13 @@ public static class ConfigMigratorV1 {

public class OTPList {
public class KeyEntry {
public string Name { get; set; }
public string Key { get; set; }
public bool IsEncrypted { get; set; }
public string Name = "";
public string Key = "";
public bool IsEncrypted = false;

}

public int Version { get; set; }
public List<KeyEntry> Entries { get; set; }
public int Version;
public List<KeyEntry> Entries = new();
}

public static void Migrate() {
Expand Down
119 changes: 77 additions & 42 deletions Main.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,14 @@
using System.Web;
using System.Windows;
using OtpNet;
using Wox.Infrastructure.Storage;
using System.Security.Cryptography;
using System.Text;
using System.Windows.Controls;

namespace Community.PowerToys.Run.Plugin.TOTP {

public class Main: IPlugin {
public class Main: IPlugin, ISavable, IReloadable, IDisposable {
public static string PluginID => "2FC51DBA9F0F42108E26602486C186C1";
private string IconCopy = "images/copy-light.png";
private string IconAdd = "images/add-light.png";
Expand All @@ -17,6 +21,30 @@ public class Main: IPlugin {
public string Name => "TOTP";

public string Description => "TOTP Code Generator";
private bool _disposed;

private PluginJsonStorage<OTPList> _storage;
private OTPList _list;

public Main() {
_storage = new PluginJsonStorage<OTPList>();
try {
ConfigMigratorV0.Migrate();
} catch (Exception) { }
try {
ConfigMigratorV1.Migrate();
} catch (Exception) { }
// Enforce PT DO not delete old version config
new StoragePowerToysVersionInfo(_storage.FilePath, 1).Close();
_list = _storage.Load();
_list.Entries.ForEach(totp => {
if (!totp.IsEncrypted) {
totp.Key = EncryptKey(totp.Key);
totp.IsEncrypted = true;
}
});
_storage.Save();
}

public List<Result> Query(Query query) {
if (query.Search.StartsWith("otpauth://totp/")) {
Expand All @@ -30,18 +58,13 @@ public List<Result> Query(Query query) {
Title = name,
SubTitle = "Add to list",
IcoPath = IconAdd,
Action = (e) => {
try {
var list = Config.LoadKeyList();
list.Add(new OTPList.KeyEntry {
Key = sercet,
Action = (e) => {
_list.Entries.Add(new OTPList.KeyEntry {
Key = EncryptKey(sercet),
Name = name,
IsEncrypted = false
IsEncrypted = true
});
Config.SaveKeyList(list);
} catch (Exception ex) {
MessageBox.Show(ex.Message + "\n" + ex.StackTrace, "PowerToys TOTP Ran into error", MessageBoxButton.OK, MessageBoxImage.Exclamation);
}
_storage.Save();
return true;
}
});
Expand All @@ -64,14 +87,12 @@ public List<Result> Query(Query query) {
var payload = queries.Get("data") ?? throw new Exception();
var decoded = Payload.Parser.ParseFrom(Convert.FromBase64String(payload));


return new List<Result> {
new Result {
Title = "Add " + decoded.OtpParameters.Count() + " items to list",
SubTitle = "From Google Authenticator App, batch " + (decoded.BatchIndex + 1).ToString() + " / " + decoded.BatchSize,
IcoPath = IconAdd,
Action = (e) => {
var list = Config.LoadKeyList();
foreach (var item in decoded.OtpParameters) {
var key = Base32Encoding.ToString(item.Secret.ToByteArray());
var name = "";
Expand All @@ -82,16 +103,13 @@ public List<Result> Query(Query query) {
} else {
name = item.Issuer + ": " + item.Name;
}
list.Add(new OTPList.KeyEntry{
_list.Entries.Add(new OTPList.KeyEntry{
Name = name,
Key = key,
IsEncrypted = false
Key = EncryptKey(key),
IsEncrypted = true
});
_storage.Save();
}
try {
Config.SaveKeyList(list);
} catch(Exception ex) {
MessageBox.Show(ex.Message + "\n" + ex.StackTrace, "PowerToys TOTP Ran into error", MessageBoxButton.OK, MessageBoxImage.Exclamation);}
return true;
}
}
Expand All @@ -109,26 +127,15 @@ public List<Result> Query(Query query) {
};
}
}
List<OTPList.KeyEntry> totpList;
try {
totpList = Config.LoadKeyList();
} catch (Exception ex) {
return new List<Result> {
new Result {
Title = ex.Message,
SubTitle = "Error when try to load config",
IcoPath = IconWarn
}
};
}

var result = new List<Result>();

totpList.ForEach(totp => {
_list.Entries.ForEach(totp => {
if (query.Search.Length != 0 && !StringMatcher.FuzzySearch(query.Search, totp.Name).Success)
return;
var key = totp.Key;
if (totp.IsEncrypted) {
key = Config.DecryptKey(key);
key = DecryptKey(key);
}
var totpInst = new Totp(Base32Encoding.ToBytes(key));
result.Add(new Result {
Expand All @@ -142,7 +149,7 @@ public List<Result> Query(Query query) {
});
});
if (result.Count == 0 && query.RawQuery.StartsWith(query.ActionKeyword)) {
if (totpList.Count == 0) {
if (_list.Entries.Count == 0) {
result.Add(new Result {
Title = "No TOTP found in config",
SubTitle = "Add TOTP to plugin by paste your setup link(otpauth://) first",
Expand All @@ -169,13 +176,6 @@ public void Init(PluginInitContext context) {
Context = context;
Context.API.ThemeChanged += OnThemeChanged;
UpdateIconPath(Context.API.GetCurrentTheme());

try {
ConfigMigratorV0.Migrate();
} catch (Exception) { }
try {
ConfigMigratorV1.Migrate();
} catch (Exception) { }
}

private void UpdateIconPath(Theme theme) {
Expand All @@ -193,5 +193,40 @@ private void UpdateIconPath(Theme theme) {
private void OnThemeChanged(Theme currentTheme, Theme newTheme) {
UpdateIconPath(newTheme);
}

public void Save() {
_storage.Save();
}

public static string DecryptKey(string encrypted) {
return Encoding.UTF8.GetString(ProtectedData.Unprotect(Convert.FromBase64String(encrypted), null, DataProtectionScope.CurrentUser));
}

public static string EncryptKey(string unencrypted) {
return Convert.ToBase64String(ProtectedData.Protect(Encoding.UTF8.GetBytes(unencrypted), null, DataProtectionScope.CurrentUser));
}

public void ReloadData() {
_list = _storage.Load();
}

public void Dispose() {
Dispose(true);
GC.SuppressFinalize(this);
}

protected virtual void Dispose(bool disposing) {
if (!_disposed && disposing) {
if (Context != null && Context.API != null) {
Context.API.ThemeChanged -= OnThemeChanged;
}

_disposed = true;
}
}

public Control CreateSettingPanel() {
throw new NotImplementedException();
}
}
}

0 comments on commit af361df

Please sign in to comment.