Skip to content

Commit 575021b

Browse files
openbulletopenbullet
committed
Add DNS block
Closes #927
1 parent a45efdc commit 575021b

9 files changed

Lines changed: 574 additions & 1 deletion

File tree

‎RuriLib.Tests/Blocks/Requests/DnsRequestBlocksTests.cs‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
using System;
2+
using RuriLib.Functions.Networking;
3+
using RuriLib.Logging;
4+
using RuriLib.Blocks.Requests.Dns;
5+
using RuriLib.Models.Bots;
6+
using RuriLib.Models.Configs;
7+
using RuriLib.Models.Data;
8+
using RuriLib.Models.Environment;
9+
using RuriLib.Tests.Utils;
10+
using System.Net;
11+
using System.Text;
12+
using System.Threading.Tasks;
13+
using Xunit;
14+
using BotProviders = RuriLib.Models.Bots.Providers;
15+
16+
namespace RuriLib.Tests.Blocks.Requests;
17+
18+
public class DnsRequestBlocksTests
19+
{
20+
[Fact]
21+
public async Task DnsLookup_DnsOverHttps_UsesCustomEndpoint()
22+
{
23+
const string payload = "{\"Answer\":[{\"data\":\"10 mx01.example.com.\"},{\"data\":\"20 mx02.example.com.\"}]}";
24+
await using var server = LocalHttpResponseServer.CreateDelayed(TimeSpan.Zero, Encoding.UTF8.GetBytes(payload),
25+
"Content-Type: application/json");
26+
var data = NewBotData();
27+
28+
var result = await Methods.LookupDnsAsync(
29+
data,
30+
"example.com",
31+
DnsRecordType.MX,
32+
DnsTransportProtocol.DnsOverHttps,
33+
server.Uri.ToString(),
34+
timeoutMilliseconds: 3000);
35+
36+
Assert.Equal(["mx01.example.com", "mx02.example.com"], result);
37+
}
38+
39+
[Fact]
40+
public async Task DnsLookup_Udp_UsesCustomServer()
41+
{
42+
await using var server = TestDnsServer.CreateARecord(IPAddress.Parse("127.0.0.42"));
43+
var data = NewBotData();
44+
45+
var result = await Methods.LookupDnsAsync(
46+
data,
47+
"example.com",
48+
DnsRecordType.A,
49+
DnsTransportProtocol.Udp,
50+
$"{server.Host}:{server.UdpPort}",
51+
timeoutMilliseconds: 3000);
52+
53+
Assert.Equal(["127.0.0.42"], result);
54+
}
55+
56+
[Fact]
57+
public async Task DnsLookup_Tcp_UsesCustomServer()
58+
{
59+
await using var server = TestDnsServer.CreateARecord(IPAddress.Parse("127.0.0.43"));
60+
var data = NewBotData();
61+
62+
var result = await Methods.LookupDnsAsync(
63+
data,
64+
"example.com",
65+
DnsRecordType.A,
66+
DnsTransportProtocol.Tcp,
67+
$"{server.Host}:{server.TcpPort}",
68+
timeoutMilliseconds: 3000);
69+
70+
Assert.Equal(["127.0.0.43"], result);
71+
}
72+
73+
private static BotData NewBotData()
74+
=> new(
75+
new BotProviders(null!),
76+
new ConfigSettings(),
77+
new BotLogger(),
78+
new DataLine("", new WordlistType()));
79+
}

‎RuriLib.Tests/Helpers/Blocks/DescriptorsRepositoryTests.cs‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,18 @@ public void GetAs_BlockIdOverride_UsesStableIdAndAsyncMethodName()
5151
Assert.True(descriptor.Async);
5252
}
5353

54+
[Fact]
55+
public void GetAs_DnsLookup_ReturnsAutoDescriptor()
56+
{
57+
var repository = new DescriptorsRepository();
58+
59+
var descriptor = repository.GetAs<AutoBlockDescriptor>("DnsLookup");
60+
61+
Assert.Equal("DnsLookup", descriptor.Id);
62+
Assert.Equal("LookupDnsAsync", descriptor.MethodName);
63+
Assert.True(descriptor.Async);
64+
}
65+
5466
[Fact]
5567
public void AsTree_ContainsAutoBlockDescriptors()
5668
{

‎RuriLib.Tests/Models/Blocks/AutoBlockInstanceTests.cs‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using RuriLib.Helpers.Blocks;
22
using RuriLib.Helpers.CSharp;
3+
using RuriLib.Functions.Networking;
34
using RuriLib.Models.Blocks;
45
using RuriLib.Models.Blocks.Settings;
56
using RuriLib.Models.Blocks.Settings.Interpolated;
@@ -222,6 +223,13 @@ public void ToSyntax_TracksExpectedPatternsAndVariableFlow()
222223
AssertSyntax(CreateTcpConnectBlock(), [],
223224
[],
224225
"await TcpConnect(data, ObjectExtensions.DynamicAsString(globals.host), 80, false, ObjectExtensions.DynamicAsInt(input.timeout)).ConfigureAwait(false);");
226+
227+
AssertSyntax(CreateDnsLookupBlock(), [],
228+
["dnsAnswers"],
229+
"LookupDnsAsync(data, ObjectExtensions.DynamicAsString(globals.query)",
230+
"\"127.0.0.1:5353\"",
231+
"1500).ConfigureAwait(false);",
232+
"data.LogVariableAssignment(nameof(dnsAnswers));");
225233
}
226234

227235
private static AutoBlockInstance CreateSubstringBlock(
@@ -300,6 +308,27 @@ private static AutoBlockInstance CreateTcpConnectBlock()
300308
return block;
301309
}
302310

311+
private static AutoBlockInstance CreateDnsLookupBlock()
312+
{
313+
var block = BlockFactory.GetBlock<AutoBlockInstance>("DnsLookup");
314+
block.OutputVariable = "dnsAnswers";
315+
316+
var query = block.Settings["query"];
317+
var recordType = block.Settings["recordType"];
318+
var transport = block.Settings["transport"];
319+
var server = block.Settings["server"];
320+
var timeout = block.Settings["timeoutMilliseconds"];
321+
322+
query.InputMode = SettingInputMode.Variable;
323+
query.InputVariableName = "globals.query";
324+
(recordType.FixedSetting as EnumSetting)!.Value = nameof(DnsRecordType.MX);
325+
(transport.FixedSetting as EnumSetting)!.Value = nameof(DnsTransportProtocol.Tcp);
326+
(server.FixedSetting as StringSetting)!.Value = "127.0.0.1:5353";
327+
(timeout.FixedSetting as IntSetting)!.Value = 1500;
328+
329+
return block;
330+
}
331+
303332
private static void AssertSyntax(
304333
BlockInstance block,
305334
List<string> definedVariables,

‎RuriLib.Tests/Utils/TestDnsServer.cs‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
using System;
2+
using System.Buffers.Binary;
3+
using System.IO;
4+
using System.Net;
5+
using System.Net.Sockets;
6+
using System.Threading;
7+
using System.Threading.Tasks;
8+
9+
namespace RuriLib.Tests.Utils;
10+
11+
internal sealed class TestDnsServer : IAsyncDisposable
12+
{
13+
private readonly UdpClient udpClient = new(new IPEndPoint(IPAddress.Loopback, 0));
14+
private readonly TcpListener tcpListener = new(IPAddress.Loopback, 0);
15+
private readonly CancellationTokenSource cancellationTokenSource = new();
16+
private readonly Task udpTask;
17+
private readonly Task tcpTask;
18+
private readonly byte[] answerBytes;
19+
20+
private TestDnsServer(byte[] answerBytes)
21+
{
22+
this.answerBytes = answerBytes;
23+
tcpListener.Start();
24+
udpTask = Task.Run(() => RunUdpAsync(cancellationTokenSource.Token), CancellationToken.None);
25+
tcpTask = Task.Run(() => RunTcpAsync(cancellationTokenSource.Token), CancellationToken.None);
26+
}
27+
28+
public string Host => "127.0.0.1";
29+
30+
public int UdpPort => ((IPEndPoint)udpClient.Client.LocalEndPoint!).Port;
31+
32+
public int TcpPort => ((IPEndPoint)tcpListener.LocalEndpoint).Port;
33+
34+
public static TestDnsServer CreateARecord(IPAddress address)
35+
=> new(address.GetAddressBytes());
36+
37+
public async ValueTask DisposeAsync()
38+
{
39+
cancellationTokenSource.Cancel();
40+
udpClient.Dispose();
41+
tcpListener.Stop();
42+
43+
try
44+
{
45+
await Task.WhenAll(udpTask, tcpTask);
46+
}
47+
catch (OperationCanceledException)
48+
{
49+
}
50+
catch (ObjectDisposedException)
51+
{
52+
}
53+
catch (SocketException)
54+
{
55+
}
56+
finally
57+
{
58+
cancellationTokenSource.Dispose();
59+
}
60+
}
61+
62+
private async Task RunUdpAsync(CancellationToken cancellationToken)
63+
{
64+
while (!cancellationToken.IsCancellationRequested)
65+
{
66+
var result = await udpClient.ReceiveAsync(cancellationToken);
67+
var response = BuildAResponse(result.Buffer, answerBytes);
68+
await udpClient.SendAsync(response, result.RemoteEndPoint, cancellationToken);
69+
}
70+
}
71+
72+
private async Task RunTcpAsync(CancellationToken cancellationToken)
73+
{
74+
while (!cancellationToken.IsCancellationRequested)
75+
{
76+
using var client = await tcpListener.AcceptTcpClientAsync(cancellationToken);
77+
await using var stream = client.GetStream();
78+
79+
var lengthBuffer = new byte[2];
80+
await ReadExactlyAsync(stream, lengthBuffer, cancellationToken);
81+
var messageLength = BinaryPrimitives.ReadUInt16BigEndian(lengthBuffer);
82+
83+
var request = new byte[messageLength];
84+
await ReadExactlyAsync(stream, request, cancellationToken);
85+
86+
var response = BuildAResponse(request, answerBytes);
87+
var prefix = new byte[2];
88+
BinaryPrimitives.WriteUInt16BigEndian(prefix, (ushort)response.Length);
89+
90+
await stream.WriteAsync(prefix, cancellationToken);
91+
await stream.WriteAsync(response, cancellationToken);
92+
}
93+
}
94+
95+
private static async Task ReadExactlyAsync(Stream stream, byte[] buffer, CancellationToken cancellationToken)
96+
{
97+
var offset = 0;
98+
99+
while (offset < buffer.Length)
100+
{
101+
var read = await stream.ReadAsync(buffer.AsMemory(offset, buffer.Length - offset), cancellationToken);
102+
if (read == 0)
103+
{
104+
throw new EndOfStreamException("Unexpected end of DNS TCP stream");
105+
}
106+
107+
offset += read;
108+
}
109+
}
110+
111+
private static byte[] BuildAResponse(byte[] request, byte[] answerBytes)
112+
{
113+
var questionLength = 12;
114+
115+
while (questionLength < request.Length && request[questionLength] != 0)
116+
{
117+
questionLength += request[questionLength] + 1;
118+
}
119+
120+
questionLength += 1 + 4;
121+
122+
using var ms = new MemoryStream();
123+
using var writer = new BinaryWriter(ms);
124+
125+
writer.Write(request[0]);
126+
writer.Write(request[1]);
127+
writer.Write((byte)0x81);
128+
writer.Write((byte)0x80);
129+
writer.Write((byte)0x00);
130+
writer.Write((byte)0x01);
131+
writer.Write((byte)0x00);
132+
writer.Write((byte)0x01);
133+
writer.Write((byte)0x00);
134+
writer.Write((byte)0x00);
135+
writer.Write((byte)0x00);
136+
writer.Write((byte)0x00);
137+
writer.Write(request, 12, questionLength - 12);
138+
writer.Write((byte)0xC0);
139+
writer.Write((byte)0x0C);
140+
writer.Write((byte)0x00);
141+
writer.Write((byte)0x01);
142+
writer.Write((byte)0x00);
143+
writer.Write((byte)0x01);
144+
writer.Write((byte)0x00);
145+
writer.Write((byte)0x00);
146+
writer.Write((byte)0x00);
147+
writer.Write((byte)0x3C);
148+
writer.Write((byte)0x00);
149+
writer.Write((byte)answerBytes.Length);
150+
writer.Write(answerBytes);
151+
152+
return ms.ToArray();
153+
}
154+
}

‎RuriLib/Blocks/Requests/Dns/Methods.cs‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
using RuriLib.Attributes;
2+
using RuriLib.Exceptions;
3+
using RuriLib.Functions.Networking;
4+
using RuriLib.Logging;
5+
using RuriLib.Models.Bots;
6+
using System.Collections.Generic;
7+
using System.Threading.Tasks;
8+
9+
namespace RuriLib.Blocks.Requests.Dns;
10+
11+
/// <summary>
12+
/// Blocks for performing DNS lookups.
13+
/// </summary>
14+
[BlockCategory("DNS", "Blocks to query DNS records", "#87ceeb")]
15+
public static class Methods
16+
{
17+
/// <summary>
18+
/// Queries DNS records for a given name.
19+
/// </summary>
20+
[Block("Queries DNS records for a given name", name = "DNS Lookup", id = "DnsLookup")]
21+
public static async Task<List<string>> LookupDnsAsync(
22+
BotData data,
23+
string query,
24+
DnsRecordType recordType = DnsRecordType.A,
25+
DnsTransportProtocol transport = DnsTransportProtocol.DnsOverHttps,
26+
[BlockParam("Server", "For UDP/TCP use host, IP or host:port. For DoH use an absolute resolve endpoint URL. Leave empty to use the default resolver.")] string server = "",
27+
int timeoutMilliseconds = 10000)
28+
{
29+
data.Logger.LogHeader();
30+
31+
List<string> answers;
32+
33+
if (transport == DnsTransportProtocol.DnsOverHttps)
34+
{
35+
answers = await RuriLib.Functions.Networking.DnsLookup
36+
.FromDnsOverHttpsAsync(
37+
query,
38+
recordType.ToString(),
39+
server,
40+
data.UseProxy ? data.Proxy : null,
41+
timeoutMilliseconds,
42+
data.CancellationToken)
43+
.ConfigureAwait(false);
44+
}
45+
else
46+
{
47+
if (data.UseProxy)
48+
{
49+
throw new BlockExecutionException("UDP and TCP DNS lookups do not support proxies");
50+
}
51+
52+
answers = await RuriLib.Functions.Networking.DnsLookup
53+
.FromNameServerAsync(
54+
query,
55+
recordType,
56+
server,
57+
transport,
58+
timeoutMilliseconds,
59+
data.CancellationToken)
60+
.ConfigureAwait(false);
61+
}
62+
63+
data.Logger.Log($"Queried {recordType} records for {query} via {transport}", LogColors.SteelBlue);
64+
data.Logger.Log(answers, LogColors.SteelBlue);
65+
66+
return answers;
67+
}
68+
}

0 commit comments

Comments
 (0)