Skip to content

Commit

Permalink
Merge remote-tracking branch 'ikbenale2/search-all'
Browse files Browse the repository at this point in the history
* ikbenale2/search-all:
  Support binary lookup mode on `search_all`
  Expose patricia_search_all on SubnetTree
  • Loading branch information
awelzel committed Jun 24, 2024
2 parents f00da03 + 34e2307 commit ebca8ed
Show file tree
Hide file tree
Showing 6 changed files with 264 additions and 55 deletions.
4 changes: 4 additions & 0 deletions README
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ A simple example which associates CIDR prefixes with strings::
Network 1
>>> print("10.20.1.1" in t)
False
>>> t.search_all("10.1.42.1")
['Network 1, Subnet 42', 'Network 1']
>>> try:
... print(t["10.20.1.1"])
... except KeyError as err:
Expand All @@ -65,6 +67,8 @@ PySubnetTree also supports IPv6 addresses and prefixes::
Company 1, Site 1
>>> t["2001:db8:fe:1234::"]
Company 1
>>> t.search_all("2001:db8:4000:abcd::1")
['Company 1, Site 1', 'Company 1']


By default, CIDR prefixes and IP addresses are given as strings.
Expand Down
62 changes: 62 additions & 0 deletions SubnetTree.cc
Original file line number Diff line number Diff line change
Expand Up @@ -300,6 +300,68 @@ PyObject* SubnetTree::lookup(int family, inx_addr addr) const
return data;
}

PyObject* SubnetTree::search_all(const char *cidr, int size) const
{
int family;
inx_addr subnet;
unsigned short mask;

if ( binary_lookup_mode ) {
if ( size == 4 ) {
family = AF_INET;
mask = 32;
}

else if ( size == 16 ) {
family = AF_INET6;
mask = 128;
}

else {
PyErr_SetString(PyExc_ValueError, "Invalid binary address. Binary addresses are 4 or 16 bytes.");
return 0;
}

memcpy(&subnet, cidr, size);
}

else if ( ! parse_cidr(cidr, &family, &subnet, &mask) ) {
PyErr_SetString(PyExc_ValueError, "Invalid CIDR.");
return 0;
}

prefix_t* sn = make_prefix();

if ( ! sn ) {
PyErr_SetString(PyExc_RuntimeError, "out of memory");
return 0;
}

bool res = set_prefix(sn, family, &subnet, mask);

if ( ! res ) {
Deref_Prefix(sn);
PyErr_SetString(PyExc_RuntimeError, "invalid subnet/prefix");
return 0;
}

patricia_node_t **outlist = nullptr;
int n;
int result = patricia_search_all(tree, sn, &outlist, &n);
Deref_Prefix(sn);

PyObject* list = PyList_New(n);
for (int i = 0; i < n; i++) {
PyObject* data = (PyObject*)outlist[i]->data;
Py_INCREF(data);
PyList_SetItem(list, i, data);
}

free(outlist);

return list;
}

PyObject* SubnetTree::prefixes(bool ipv4_native /*=false*/, bool with_len /*=true*/) const
{
char buf[INET6_ADDRSTRLEN];
Expand Down
3 changes: 3 additions & 0 deletions SubnetTree.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,9 @@ def remove(self, *args):
def lookup(self, *args):
return _SubnetTree.SubnetTree_lookup(self, *args)

def search_all(self, cidr):
return _SubnetTree.SubnetTree_search_all(self, cidr)

def prefixes(self, ipv4_native=False, with_len=True):
return _SubnetTree.SubnetTree_prefixes(self, ipv4_native, with_len)

Expand Down
190 changes: 141 additions & 49 deletions SubnetTree_wrap.cc
Original file line number Diff line number Diff line change
Expand Up @@ -3459,29 +3459,13 @@ SWIG_AsVal_unsigned_SS_short (PyObject * obj, unsigned short *val)
}


SWIGINTERN int
SWIG_AsVal_int (PyObject * obj, int *val)
{
long v;
int res = SWIG_AsVal_long (obj, &v);
if (SWIG_IsOK(res)) {
if ((v < INT_MIN || v > INT_MAX)) {
return SWIG_OverflowError;
} else {
if (val) *val = static_cast< int >(v);
}
}
return res;
}


SWIGINTERNINLINE PyObject*
SWIG_From_bool (bool value)
{
return PyBool_FromLong(value ? 1 : 0);
}

SWIGINTERN PyObject *SubnetTree___contains____SWIG_0(SubnetTree *self,char *cidr,int size){
SWIGINTERN PyObject *SubnetTree___contains____SWIG_0(SubnetTree *self,char const *cidr,int size){
if ( ! cidr ) {
PyErr_SetString(PyExc_TypeError, "index must be string");
return 0;
Expand Down Expand Up @@ -3517,7 +3501,7 @@ SWIGINTERN PyObject *SubnetTree___contains____SWIG_1(SubnetTree *self,unsigned l
else
Py_RETURN_FALSE;
}
SWIGINTERN PyObject *SubnetTree___getitem__(SubnetTree *self,char *cidr,int size){
SWIGINTERN PyObject *SubnetTree___getitem__(SubnetTree *self,char const *cidr,int size){
if ( ! cidr ) {
PyErr_SetString(PyExc_TypeError, "index must be string");
return 0;
Expand Down Expand Up @@ -4197,38 +4181,70 @@ SWIGINTERN PyObject *_wrap_SubnetTree_lookup__SWIG_0(PyObject *SWIGUNUSEDPARM(se
int arg3 ;
void *argp1 = 0 ;
int res1 = 0 ;
int res2 ;
char *buf2 = 0 ;
int alloc2 = 0 ;
int val3 ;
int ecode3 = 0 ;
PyObject *ascii2 ;
PyObject * obj0 = 0 ;
PyObject * obj1 = 0 ;
PyObject * obj2 = 0 ;
PyObject *result = 0 ;

if (!PyArg_ParseTuple(args,(char *)"OOO:SubnetTree_lookup",&obj0,&obj1,&obj2)) SWIG_fail;
{
ascii2 = NULL;
}
if (!PyArg_ParseTuple(args,(char *)"OO:SubnetTree_lookup",&obj0,&obj1)) SWIG_fail;
res1 = SWIG_ConvertPtr(obj0, &argp1,SWIGTYPE_p_SubnetTree, 0 | 0 );
if (!SWIG_IsOK(res1)) {
SWIG_exception_fail(SWIG_ArgError(res1), "in method '" "SubnetTree_lookup" "', argument " "1"" of type '" "SubnetTree const *""'");
}
arg1 = reinterpret_cast< SubnetTree * >(argp1);
res2 = SWIG_AsCharPtrAndSize(obj1, &buf2, NULL, &alloc2);
if (!SWIG_IsOK(res2)) {
SWIG_exception_fail(SWIG_ArgError(res2), "in method '" "SubnetTree_lookup" "', argument " "2"" of type '" "char const *""'");

Py_ssize_t len;

#if PY_MAJOR_VERSION >= 3
if ( PyUnicode_Check(obj1) )
{
ascii2 = PyUnicode_AsASCIIString(obj1);

if ( ! ascii2 )
{
PyErr_SetString(PyExc_TypeError, "Expected a ASCII encodable string or bytes");
return NULL;
}

PyBytes_AsStringAndSize(ascii2, &arg2, &len);
arg3 = len;
}
arg2 = reinterpret_cast< char * >(buf2);
ecode3 = SWIG_AsVal_int(obj2, &val3);
if (!SWIG_IsOK(ecode3)) {
SWIG_exception_fail(SWIG_ArgError(ecode3), "in method '" "SubnetTree_lookup" "', argument " "3"" of type '" "int""'");
}
arg3 = static_cast< int >(val3);
else if ( PyBytes_Check(obj1) )
{
PyBytes_AsStringAndSize(obj1, &arg2, &len);
arg3 = len;
}
else
{
PyErr_SetString(PyExc_TypeError, "Expected a string or bytes");
return NULL;
}
#else
if ( ! PyString_Check(obj1) )
{
PyErr_SetString(PyExc_TypeError, "Expected a string or bytes");
return NULL;
}
else
{
PyString_AsStringAndSize(obj1, &arg2, &len);
arg3 = len;
}
#endif

result = (PyObject *)((SubnetTree const *)arg1)->lookup((char const *)arg2,arg3);
resultobj = result;
if (alloc2 == SWIG_NEWOBJ) delete[] buf2;
{
Py_XDECREF(ascii2);
}
return resultobj;
fail:
if (alloc2 == SWIG_NEWOBJ) delete[] buf2;
{
Py_XDECREF(ascii2);
}
return NULL;
}

Expand Down Expand Up @@ -4266,14 +4282,14 @@ SWIGINTERN PyObject *_wrap_SubnetTree_lookup__SWIG_1(PyObject *SWIGUNUSEDPARM(se

SWIGINTERN PyObject *_wrap_SubnetTree_lookup(PyObject *self, PyObject *args) {
Py_ssize_t argc;
PyObject *argv[4] = {
PyObject *argv[3] = {
0
};
Py_ssize_t ii;

if (!PyTuple_Check(args)) SWIG_fail;
argc = args ? PyObject_Length(args) : 0;
for (ii = 0; (ii < 3) && (ii < argc); ii++) {
for (ii = 0; (ii < 2) && (ii < argc); ii++) {
argv[ii] = PyTuple_GET_ITEM(args,ii);
}
if (argc == 2) {
Expand All @@ -4291,22 +4307,22 @@ SWIGINTERN PyObject *_wrap_SubnetTree_lookup(PyObject *self, PyObject *args) {
}
}
}
if (argc == 3) {
if (argc == 2) {
int _v;
void *vptr = 0;
int res = SWIG_ConvertPtr(argv[0], &vptr, SWIGTYPE_p_SubnetTree, 0);
_v = SWIG_CheckState(res);
if (_v) {
int res = SWIG_AsCharPtrAndSize(argv[1], 0, NULL, 0);
_v = SWIG_CheckState(res);
{
// The typemap above will check types and throw a type error when
// needed, so just let everything through.
_v = 1;
}
if (_v) {
{
int res = SWIG_AsVal_int(argv[2], NULL);
_v = SWIG_CheckState(res);
}
if (_v) {
if (argc <= 2) {
return _wrap_SubnetTree_lookup__SWIG_0(self, args);
}
return _wrap_SubnetTree_lookup__SWIG_0(self, args);
}
}
}
Expand All @@ -4320,6 +4336,81 @@ SWIGINTERN PyObject *_wrap_SubnetTree_lookup(PyObject *self, PyObject *args) {
}


SWIGINTERN PyObject *_wrap_SubnetTree_search_all(PyObject *SWIGUNUSEDPARM(self), PyObject *args) {
PyObject *resultobj = 0;
SubnetTree *arg1 = (SubnetTree *) 0 ;
char *arg2 = (char *) 0 ;
int arg3 ;
void *argp1 = 0 ;
int res1 = 0 ;
PyObject *ascii2 ;
PyObject * obj0 = 0 ;
PyObject * obj1 = 0 ;
PyObject *result = 0 ;

{
ascii2 = NULL;
}
if (!PyArg_ParseTuple(args,(char *)"OO:SubnetTree_search_all",&obj0,&obj1)) SWIG_fail;
res1 = SWIG_ConvertPtr(obj0, &argp1,SWIGTYPE_p_SubnetTree, 0 | 0 );
if (!SWIG_IsOK(res1)) {
SWIG_exception_fail(SWIG_ArgError(res1), "in method '" "SubnetTree_search_all" "', argument " "1"" of type '" "SubnetTree const *""'");
}
arg1 = reinterpret_cast< SubnetTree * >(argp1);

Py_ssize_t len;

#if PY_MAJOR_VERSION >= 3
if ( PyUnicode_Check(obj1) )
{
ascii2 = PyUnicode_AsASCIIString(obj1);

if ( ! ascii2 )
{
PyErr_SetString(PyExc_TypeError, "Expected a ASCII encodable string or bytes");
return NULL;
}

PyBytes_AsStringAndSize(ascii2, &arg2, &len);
arg3 = len;
}
else if ( PyBytes_Check(obj1) )
{
PyBytes_AsStringAndSize(obj1, &arg2, &len);
arg3 = len;
}
else
{
PyErr_SetString(PyExc_TypeError, "Expected a string or bytes");
return NULL;
}
#else
if ( ! PyString_Check(obj1) )
{
PyErr_SetString(PyExc_TypeError, "Expected a string or bytes");
return NULL;
}
else
{
PyString_AsStringAndSize(obj1, &arg2, &len);
arg3 = len;
}
#endif

result = (PyObject *)((SubnetTree const *)arg1)->search_all((char const *)arg2,arg3);
resultobj = result;
{
Py_XDECREF(ascii2);
}
return resultobj;
fail:
{
Py_XDECREF(ascii2);
}
return NULL;
}


SWIGINTERN PyObject *_wrap_SubnetTree_prefixes__SWIG_0(PyObject *SWIGUNUSEDPARM(self), PyObject *args) {
PyObject *resultobj = 0;
SubnetTree *arg1 = (SubnetTree *) 0 ;
Expand Down Expand Up @@ -4661,7 +4752,7 @@ SWIGINTERN PyObject *_wrap_SubnetTree___contains____SWIG_0(PyObject *SWIGUNUSEDP
}
#endif

result = (PyObject *)SubnetTree___contains____SWIG_0(arg1,arg2,arg3);
result = (PyObject *)SubnetTree___contains____SWIG_0(arg1,(char const *)arg2,arg3);
resultobj = result;
{
Py_XDECREF(ascii2);
Expand Down Expand Up @@ -4756,7 +4847,7 @@ SWIGINTERN PyObject *_wrap_SubnetTree___contains__(PyObject *self, PyObject *arg
fail:
SWIG_SetErrorMsg(PyExc_NotImplementedError,"Wrong number or type of arguments for overloaded function 'SubnetTree___contains__'.\n"
" Possible C/C++ prototypes are:\n"
" SubnetTree::__contains__(char *,int)\n"
" SubnetTree::__contains__(char const *,int)\n"
" SubnetTree::__contains__(unsigned long)\n");
return 0;
}
Expand Down Expand Up @@ -4823,7 +4914,7 @@ SWIGINTERN PyObject *_wrap_SubnetTree___getitem__(PyObject *SWIGUNUSEDPARM(self)
}
#endif

result = (PyObject *)SubnetTree___getitem__(arg1,arg2,arg3);
result = (PyObject *)SubnetTree___getitem__(arg1,(char const *)arg2,arg3);
resultobj = result;
{
Py_XDECREF(ascii2);
Expand Down Expand Up @@ -4929,6 +5020,7 @@ static PyMethodDef SwigMethods[] = {
{ (char *)"SubnetTree_insert", _wrap_SubnetTree_insert, METH_VARARGS, NULL},
{ (char *)"SubnetTree_remove", _wrap_SubnetTree_remove, METH_VARARGS, NULL},
{ (char *)"SubnetTree_lookup", _wrap_SubnetTree_lookup, METH_VARARGS, NULL},
{ (char *)"SubnetTree_search_all", _wrap_SubnetTree_search_all, METH_VARARGS, NULL},
{ (char *)"SubnetTree_prefixes", _wrap_SubnetTree_prefixes, METH_VARARGS, NULL},
{ (char *)"SubnetTree_get_binary_lookup_mode", _wrap_SubnetTree_get_binary_lookup_mode, METH_VARARGS, NULL},
{ (char *)"SubnetTree_set_binary_lookup_mode", _wrap_SubnetTree_set_binary_lookup_mode, METH_VARARGS, NULL},
Expand Down
Loading

0 comments on commit ebca8ed

Please sign in to comment.