drastic nsfw and safebooru improvements

src/NadekoBot/DataStructures/SearchImageCacher.cs

@@ -0,0 +1,213 @@
+using NadekoBot.Extensions;
+using NadekoBot.Services;
+using Newtonsoft.Json;
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+using System.Xml;
+using System.Xml.Linq;
+
+namespace NadekoBot.DataStructures
+{
+    public class SearchImageCacher
+    {
+        private readonly NadekoRandom _rng;
+        private readonly ConcurrentDictionary<DapiSearchType, SemaphoreSlim> _locks = new ConcurrentDictionary<DapiSearchType, SemaphoreSlim>();
+
+        private readonly SortedSet<ImageCacherObject> _cache;
+
+        public SearchImageCacher()
+        {
+            _rng = new NadekoRandom();
+            _cache = new SortedSet<ImageCacherObject>();
+        }
+
+        public async Task<ImageCacherObject> GetImage(string tag, bool forceExplicit, DapiSearchType type)
+        {
+            tag = tag?.ToLowerInvariant();
+
+            if (type == DapiSearchType.E621)
+                tag = tag?.Replace("yuri", "female/female");
+
+            var _lock = GetLock(type);
+            await _lock.WaitAsync();
+            try
+            {
+                ImageCacherObject[] imgs;
+                if (!string.IsNullOrWhiteSpace(tag))
+                {
+                    imgs = _cache.Where(x => x.Tags.IsSupersetOf(tag.Split('+')) && x.SearchType == type && (!forceExplicit || x.Rating == "e")).ToArray();
+                }
+                else
+                {
+                    tag = null;
+                    imgs = _cache.Where(x => x.SearchType == type).ToArray();
+                }
+                ImageCacherObject img;
+                if (imgs.Length == 0)
+                    img = null;
+                else
+                    img = imgs[_rng.Next(imgs.Length)];
+
+                if (img != null)
+                {
+                    _cache.Remove(img);
+                    return img;
+                }
+                else
+                {
+                    var images = await DownloadImages(tag, forceExplicit, type).ConfigureAwait(false);
+                    if (images.Length == 0)
+                        return null;
+                    var toReturn = images[_rng.Next(images.Length)];
+                    foreach (var dledImg in images)
+                    {
+                        if(dledImg != toReturn)
+                            _cache.Add(dledImg);
+                    }
+                    return toReturn;
+                }                    
+            }
+            finally
+            {
+                _lock.Release();
+            }
+        }
+
+        private SemaphoreSlim GetLock(DapiSearchType type)
+        {
+            return _locks.GetOrAdd(type, _ => new SemaphoreSlim(1, 1));
+        }
+
+        public async Task<ImageCacherObject[]> DownloadImages(string tag, bool isExplicit, DapiSearchType type)
+        {
+            Console.WriteLine($"Loading extra images from {type}");
+            tag = tag?.Replace(" ", "_").ToLowerInvariant();
+            if (isExplicit)
+                tag = "rating%3Aexplicit+" + tag;
+            var website = "";
+            switch (type)
+            {
+                case DapiSearchType.Safebooru:
+                    website = $"https://safebooru.org/index.php?page=dapi&s=post&q=index&limit=1000&tags={tag}";
+                    break;
+                case DapiSearchType.E621:
+                    website = $"https://e621.net/post/index.json?limit=1000&tags={tag}";
+                    break;
+                case DapiSearchType.Danbooru:
+                    website = $"https://danbooru.donmai.us/posts.json?limit=200&tags={tag}";
+                    break;
+                case DapiSearchType.Gelbooru:
+                    website = $"http://gelbooru.com/index.php?page=dapi&s=post&q=index&limit=1000&tags={tag}";
+                    break;
+                case DapiSearchType.Rule34:
+                    website = $"https://rule34.xxx/index.php?page=dapi&s=post&q=index&limit=100&tags={tag}";
+                    break;
+                case DapiSearchType.Konachan:
+                    website = $"https://konachan.com/post.json?s=post&q=index&limit=1000&tags={tag}";
+                    break;
+                case DapiSearchType.Yandere:
+                    website = $"https://yande.re/post.json?limit=1000&tags={tag}";
+                    break;
+            }
+
+            using (var http = new HttpClient())
+            {
+                http.AddFakeHeaders();
+                
+                if (type == DapiSearchType.Konachan || type == DapiSearchType.Yandere || 
+                    type == DapiSearchType.E621 || type == DapiSearchType.Danbooru)
+                {
+                    var data = await http.GetStringAsync(website).ConfigureAwait(false);
+                    return JsonConvert.DeserializeObject<DapiImageObject[]>(data)
+                        .Where(x => x.File_Url != null)
+                        .Select(x => new ImageCacherObject(x, type))
+                        .ToArray();
+                }
+
+                return (await LoadXmlAsync(website, type)).ToArray();
+            }
+        }
+
+        private async Task<ImageCacherObject[]> LoadXmlAsync(string website, DapiSearchType type)
+        {
+            var list = new List<ImageCacherObject>(1000);
+            using (var http = new HttpClient())
+            {
+                using (var reader = XmlReader.Create(await http.GetStreamAsync(website), new XmlReaderSettings()
+                {
+                    Async = true,
+                }))
+                {
+                    while (await reader.ReadAsync())
+                    {
+                        if (reader.NodeType == XmlNodeType.Element &&
+                            reader.Name == "post")
+                        {
+                            list.Add(new ImageCacherObject(new DapiImageObject()
+                            {
+                                File_Url = reader["file_url"],
+                                Tags = reader["tags"],
+                                Rating = reader["rating"] ?? "e"
+                               
+                            }, type));
+                        }
+                    }
+                }
+            }
+            return list.ToArray();
+        }
+    }
+
+    public class ImageCacherObject : IComparable<ImageCacherObject>
+    {
+        public DapiSearchType SearchType { get; }
+        public string FileUrl { get; }
+        public HashSet<string> Tags { get; }
+        public string Rating { get; }
+
+        public ImageCacherObject(DapiImageObject obj, DapiSearchType type)
+        {
+            if (type == DapiSearchType.Danbooru)
+                this.FileUrl = "https://danbooru.donmai.us" + obj.File_Url;
+            else
+                this.FileUrl = obj.File_Url.StartsWith("http") ? obj.File_Url : "https:" + obj.File_Url;
+            this.SearchType = type;
+            this.Rating = obj.Rating;
+            this.Tags = new HashSet<string>((obj.Tags ?? obj.Tag_String).Split(' '));
+        }
+
+        public override string ToString()
+        {
+            return FileUrl;
+        }
+
+        public int CompareTo(ImageCacherObject other)
+        {
+            return FileUrl.CompareTo(other.FileUrl);
+        }
+    }
+
+    public class DapiImageObject
+    {
+        public string File_Url { get; set; }
+        public string Tags { get; set; }
+        public string Tag_String { get; set; }
+        public string Rating { get; set; }
+    }
+
+    public enum DapiSearchType
+    {
+        Safebooru,
+        E621,
+        Gelbooru,
+        Konachan,
+        Rule34,
+        Yandere,
+        Danbooru
+    }
+}

src/NadekoBot/Modules/NSFW/NSFW.cs

@@ -12,6 +12,7 @@ using System.Xml;
 using System.Threading;
 using System.Collections.Concurrent;
 using NadekoBot.Services.Searches;
+using NadekoBot.DataStructures;
 
 namespace NadekoBot.Modules.NSFW
 {
@@ -28,29 +29,12 @@ namespace NadekoBot.Modules.NSFW
 
         private async Task InternalHentai(IMessageChannel channel, string tag, bool noError)
         {
-            tag = tag?.Trim() ?? "";
-
-            tag = "rating%3Aexplicit+" + tag;
-
             var rng = new NadekoRandom();
-            var provider = Task.FromResult("");
-            switch (rng.Next(0, 4))
-            {
-                case 0:
-                    provider = GetDanbooruImageLink(tag);
-                    break;
-                case 1:
-                    provider = GetGelbooruImageLink(tag);
-                    break;
-                case 2:
-                    provider = GetKonachanImageLink(tag);
-                    break;
-                case 3:
-                    provider = GetYandereImageLink(tag);
-                    break;
-            }
-            var link = await provider.ConfigureAwait(false);
-            if (string.IsNullOrWhiteSpace(link))
+            var arr = Enum.GetValues(typeof(DapiSearchType));
+            var type = (DapiSearchType)arr.GetValue(new NadekoRandom().Next(2, arr.Length));
+            var img = await _service.DapiSearch(tag, type, Context.Guild?.Id, true).ConfigureAwait(false);
+
+            if (img == null)
             {
                 if (!noError)
                     await ReplyErrorLocalized("not_found").ConfigureAwait(false);
@@ -58,8 +42,8 @@ namespace NadekoBot.Modules.NSFW
             }
 
             await channel.EmbedAsync(new EmbedBuilder().WithOkColor()
-                .WithImageUrl(link)
-                .WithDescription($"[{GetText("tag")}: {tag}]({link})"))
+                .WithImageUrl(img.FileUrl)
+                .WithDescription($"[{GetText("tag")}: {tag}]({img})"))
                 .ConfigureAwait(false);
         }
 
@@ -112,110 +96,58 @@ namespace NadekoBot.Modules.NSFW
                 interval,
                 string.Join(", ", tagsArr)).ConfigureAwait(false);
         }
-
+#endif
 
         [NadekoCommand, Usage, Description, Aliases]
         public async Task HentaiBomb([Remainder] string tag = null)
         {
-            if (!_hentaiBombBlacklist.Add(Context.User.Id))
+            if (!_hentaiBombBlacklist.Add(Context.Guild?.Id ?? Context.User.Id))
                 return;
             try
             {
-                tag = tag?.Trim() ?? "";
-                tag = "rating%3Aexplicit+" + tag;
+                var images = await Task.WhenAll(_service.DapiSearch(tag, DapiSearchType.Gelbooru, Context.Guild?.Id, true),
+                                                _service.DapiSearch(tag, DapiSearchType.Danbooru, Context.Guild?.Id, true),
+                                                _service.DapiSearch(tag, DapiSearchType.Konachan, Context.Guild?.Id, true),
+                                                _service.DapiSearch(tag, DapiSearchType.Yandere, Context.Guild?.Id, true)).ConfigureAwait(false);
 
-                var links = await Task.WhenAll(GetGelbooruImageLink(tag),
-                                               GetDanbooruImageLink(tag),
-                                               GetKonachanImageLink(tag),
-                                               GetYandereImageLink(tag)).ConfigureAwait(false);
-
-                var linksEnum = links?.Where(l => l != null).ToArray();
-                if (links == null || !linksEnum.Any())
+                var linksEnum = images?.Where(l => l != null).ToArray();
+                if (images == null || !linksEnum.Any())
                 {
                     await ReplyErrorLocalized("not_found").ConfigureAwait(false);
                     return;
                 }
 
-                await Context.Channel.SendMessageAsync(string.Join("\n\n", linksEnum)).ConfigureAwait(false);
+                await Context.Channel.SendMessageAsync(string.Join("\n\n", linksEnum.Select(x => x.FileUrl))).ConfigureAwait(false);
             }
             finally
             {
-                await Task.Delay(5000).ConfigureAwait(false);
-                _hentaiBombBlacklist.TryRemove(Context.User.Id);
+                _hentaiBombBlacklist.TryRemove(Context.Guild?.Id ?? Context.User.Id);
             }
         }
-#endif
+
         [NadekoCommand, Usage, Description, Aliases]
         public Task Yandere([Remainder] string tag = null)
-            => InternalDapiCommand(tag, DapiSearchType.Yandere);
+            => InternalDapiCommand(tag, DapiSearchType.Yandere, false);
 
         [NadekoCommand, Usage, Description, Aliases]
         public Task Konachan([Remainder] string tag = null)
-            => InternalDapiCommand(tag, DapiSearchType.Konachan);
+            => InternalDapiCommand(tag, DapiSearchType.Konachan, false);
 
         [NadekoCommand, Usage, Description, Aliases]
-        public async Task E621([Remainder] string tag = null)
-        {
-            tag = tag?.Trim() ?? "";
-
-            var url = await GetE621ImageLink(tag).ConfigureAwait(false);
-
-            if (url == null)
-                await ReplyErrorLocalized("not_found").ConfigureAwait(false);
-            else
-                await Context.Channel.EmbedAsync(new EmbedBuilder().WithOkColor()
-                    .WithDescription(Context.User.Mention + " " + tag)
-                    .WithImageUrl(url)
-                    .WithFooter(efb => efb.WithText("e621")))
-                    .ConfigureAwait(false);
-        }
+        public Task E621([Remainder] string tag = null)
+            => InternalDapiCommand(tag, DapiSearchType.E621, false);
 
         [NadekoCommand, Usage, Description, Aliases]
         public Task Rule34([Remainder] string tag = null)
-            => InternalDapiCommand(tag, DapiSearchType.Rule34);
+            => InternalDapiCommand(tag, DapiSearchType.Rule34, false);
 
         [NadekoCommand, Usage, Description, Aliases]
-        public async Task Danbooru([Remainder] string tag = null)
-        {
-            tag = tag?.Trim() ?? "";
-
-            var url = await GetDanbooruImageLink(tag).ConfigureAwait(false);
-
-            if (url == null)
-                await ReplyErrorLocalized("not_found").ConfigureAwait(false);
-            else
-                await Context.Channel.EmbedAsync(new EmbedBuilder().WithOkColor()
-                    .WithDescription(Context.User.Mention + " " + tag)
-                    .WithImageUrl(url)
-                    .WithFooter(efb => efb.WithText("Danbooru")))
-                    .ConfigureAwait(false);
-        }
-
-        public static Task<string> GetDanbooruImageLink(string tag) => Task.Run(async () =>
-        {
-            try
-            {
-                using (var http = new HttpClient())
-                {
-                    http.AddFakeHeaders();
-                    var data = await http.GetStreamAsync("https://danbooru.donmai.us/posts.xml?limit=100&tags=" + tag).ConfigureAwait(false);
-                    var doc = new XmlDocument();
-                    doc.Load(data);
-                    var nodes = doc.GetElementsByTagName("file-url");
-
-                    var node = nodes[new NadekoRandom().Next(0, nodes.Count)];
-                    return "https://danbooru.donmai.us" + node.InnerText;
-                }
-            }
-            catch
-            {
-                return null;
-            }
-        });
+        public Task Danbooru([Remainder] string tag = null)
+            => InternalDapiCommand(tag, DapiSearchType.Danbooru, false);
 
         [NadekoCommand, Usage, Description, Aliases]
         public Task Gelbooru([Remainder] string tag = null)
-            => InternalDapiCommand(tag, DapiSearchType.Gelbooru);
+            => InternalDapiCommand(tag, DapiSearchType.Gelbooru, false);
 
         [NadekoCommand, Usage, Description, Aliases]
         public async Task Boobs()
@@ -253,52 +185,16 @@ namespace NadekoBot.Modules.NSFW
             }
         }
 
-        public static Task<string> GetE621ImageLink(string tag) => Task.Run(async () =>
+        public async Task InternalDapiCommand(string tag, DapiSearchType type, bool forceExplicit)
         {
-            try
-            {
-                using (var http = new HttpClient())
-                {
-                    http.AddFakeHeaders();
-                    var data = await http.GetStreamAsync("http://e621.net/post/index.xml?tags=" + tag).ConfigureAwait(false);
-                    var doc = new XmlDocument();
-                    doc.Load(data);
-                    var nodes = doc.GetElementsByTagName("file_url");
+            var imgObj = await _service.DapiSearch(tag, type, Context.Guild?.Id, forceExplicit).ConfigureAwait(false);
 
-                    var node = nodes[new NadekoRandom().Next(0, nodes.Count)];
-                    return node.InnerText;
-                }
-            }
-            catch
-            {
-                return null;
-            }
-        });
-
-        public Task<string> GetRule34ImageLink(string tag) =>
-            _service.DapiSearch(tag, DapiSearchType.Rule34);
-
-        public Task<string> GetYandereImageLink(string tag) =>
-            _service.DapiSearch(tag, DapiSearchType.Yandere);
-
-        public Task<string> GetKonachanImageLink(string tag) =>
-            _service.DapiSearch(tag, DapiSearchType.Konachan);
-
-        public Task<string> GetGelbooruImageLink(string tag) =>
-            _service.DapiSearch(tag, DapiSearchType.Gelbooru);
-
-        public async Task InternalDapiCommand(string tag, DapiSearchType type)
-        {
-            tag = tag?.Trim() ?? "";
-
-            var url = await _service.DapiSearch(tag, type).ConfigureAwait(false);
-
-            if (url == null)
+            if (imgObj == null)
                 await ReplyErrorLocalized("not_found").ConfigureAwait(false);
             else
                 await Context.Channel.EmbedAsync(new EmbedBuilder().WithOkColor()
-                    .WithDescription($"{Context.User} [{tag}]({url}) ")
-                    .WithImageUrl(url)
+                    .WithDescription($"{Context.User} [{tag ?? "url"}]({imgObj}) ")
+                    .WithImageUrl(imgObj.FileUrl)
                     .WithFooter(efb => efb.WithText(type.ToString()))).ConfigureAwait(false);
         }
     }

src/NadekoBot/Modules/Searches/Searches.cs

@@ -21,6 +21,7 @@ using NadekoBot.Attributes;
 using Discord.Commands;
 using ImageSharp;
 using NadekoBot.Services.Searches;
+using NadekoBot.DataStructures;
 
 namespace NadekoBot.Modules.Searches
 {
@@ -791,14 +792,14 @@ namespace NadekoBot.Modules.Searches
 
             tag = tag?.Trim() ?? "";
 
-            var url = await _searches.DapiSearch(tag, type).ConfigureAwait(false);
+            var imgObj = await _searches.DapiSearch(tag, type, Context.Guild?.Id).ConfigureAwait(false);
 
-            if (url == null)
+            if (imgObj == null)
                 await channel.SendErrorAsync(umsg.Author.Mention + " " + GetText("no_results"));
             else
                 await channel.EmbedAsync(new EmbedBuilder().WithOkColor()
-                    .WithDescription($"{umsg.Author.Mention} [{tag}]({url})")
-                    .WithImageUrl(url)
+                    .WithDescription($"{umsg.Author.Mention} [{tag ?? "url"}]({imgObj.FileUrl})")
+                    .WithImageUrl(imgObj.FileUrl)
                     .WithFooter(efb => efb.WithText(type.ToString()))).ConfigureAwait(false);
         }
 

src/NadekoBot/Services/CustomReactions/Extensions.cs

@@ -41,7 +41,7 @@ namespace NadekoBot.Services.CustomReactions
                 if (img?.Source == null)
                     return "";
 
-                return " "+img.Source.Replace("b.", ".") + " ";
+                return " " + img.Source.Replace("b.", ".") + " ";
             } }
         };
 

src/NadekoBot/Services/Impl/StatsService.cs

@@ -17,7 +17,7 @@ namespace NadekoBot.Services.Impl
         private readonly IBotCredentials _creds;
         private readonly DateTime _started;
 
-        public const string BotVersion = "1.51";
+        public const string BotVersion = "1.52";
 
         public string Author => "Kwoth#2560";
         public string Library => "Discord.Net";

src/NadekoBot/Services/Searches/SearchesService.cs

@@ -1,11 +1,13 @@
 using Discord;
 using Discord.WebSocket;
+using NadekoBot.DataStructures;
 using NadekoBot.Extensions;
 using Newtonsoft.Json;
 using NLog;
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Diagnostics;
 using System.IO;
 using System.Net.Http;
 using System.Threading.Tasks;
@@ -31,6 +33,8 @@ namespace NadekoBot.Services.Searches
         public List<WoWJoke> WowJokes { get; } = new List<WoWJoke>();
         public List<MagicItem> MagicItems { get; } = new List<MagicItem>();
 
+        private readonly ConcurrentDictionary<ulong?, SearchImageCacher> _imageCacher = new ConcurrentDictionary<ulong?, SearchImageCacher>();
+
         public SearchesService(DiscordSocketClient client, IGoogleApiService google, DbService db)
         {
             _client = client;
@@ -113,63 +117,12 @@ namespace NadekoBot.Services.Searches
             return (await _google.Translate(text, from, to).ConfigureAwait(false)).SanitizeMentions();
         }
 
-        public async Task<string> DapiSearch(string tag, DapiSearchType type)
+        public Task<ImageCacherObject> DapiSearch(string tag, DapiSearchType type, ulong? guild, bool isExplicit = false)
         {
-            tag = tag?.Replace(" ", "_");
-            var website = "";
-            switch (type)
-            {
-                case DapiSearchType.Safebooru:
-                    website = $"https://safebooru.org/index.php?page=dapi&s=post&q=index&limit=100&tags={tag}";
-                    break;
-                case DapiSearchType.Gelbooru:
-                    website = $"http://gelbooru.com/index.php?page=dapi&s=post&q=index&limit=100&tags={tag}";
-                    break;
-                case DapiSearchType.Rule34:
-                    website = $"https://rule34.xxx/index.php?page=dapi&s=post&q=index&limit=100&tags={tag}";
-                    break;
-                case DapiSearchType.Konachan:
-                    website = $"https://konachan.com/post.xml?s=post&q=index&limit=100&tags={tag}";
-                    break;
-                case DapiSearchType.Yandere:
-                    website = $"https://yande.re/post.xml?limit=100&tags={tag}";
-                    break;
-            }
-            try
-            {
-                var toReturn = await Task.Run(async () =>
-                {
-                    using (var http = new HttpClient())
-                    {
-                        http.AddFakeHeaders();
-                        var data = await http.GetStreamAsync(website).ConfigureAwait(false);
-                        var doc = new XmlDocument();
-                        doc.Load(data);
+            var cacher = _imageCacher.GetOrAdd(guild, (key) => new SearchImageCacher());
             
-                        var node = doc.LastChild.ChildNodes[new NadekoRandom().Next(0, doc.LastChild.ChildNodes.Count)];
-
-                        var url = node.Attributes["file_url"].Value;
-                        if (!url.StartsWith("http"))
-                            url = "https:" + url;
-                        return url;
+            return cacher.GetImage(tag, isExplicit, type);
         }
-                }).ConfigureAwait(false);
-                return toReturn;
-            }
-            catch
-            {
-                return null;
-            }
-        }
-    }
-
-    public enum DapiSearchType
-    {
-        Safebooru,
-        Gelbooru,
-        Konachan,
-        Rule34,
-        Yandere
     }
     
     public struct UserChannelPair
@@ -178,7 +131,6 @@ namespace NadekoBot.Services.Searches
         public ulong ChannelId { get; set; }
     }
 
-
     public class StreamStatus
     {
         public bool IsLive { get; set; }